%% Fitting the motion referenced model with the update at the error angle location to the data
clearvars
currentFolder = pwd;
OVERWRITE = 0;                                                              % redo the fitting procedure
addpath(genpath('Matlab_functions/bads-master')) 
addpath(genpath('Matlab_functions')) 
path = currentFolder(1:end-8); 
data_folder = 'data\processed\';
fits_folder = 'data\fits\';
%% what are all possible target angles
load(sprintf('%s%ss%03d%s',path,data_folder,1,'.mat'))
aPos = unique(subject.TA);
% 
CW = [-1,-1,-1,1,1,1,-1,1,1,-1,-1,1,-1,1,-1,1,1,1,-1,-1,-1,1,1,-1];         % -1 no/CCW-CW; 1 yes/CW-CCW 

%% load in error angles
load(sprintf('%s%sUpdateAngle',path,data_folder),'ThetaInDeg')

%% what type of model to fit:
model = 'plan_based';
%% trialtypes 
%   Forces
TT.ff = [0 1 2] ;                                                           
% 0: no force trials,                
% 1: first force field trials,          
% 2: second force field trials
%   Clamp Trials
TT.chan =[3 4 5 6 7];                                                       
% 5: clamp trials in no force field, 
% 3: clamp trials in first force field,
% 6: clamp trials in second force field,
% 4: spontaneous recovery, 
% 7: post break clamp trials
TT.fit = [3 4 6];
% trials used during fitting

%% motor primitives 
P.sigma = 30;
P.theta =360;
P.nj = 100;
[~,g] = MotorPrimitives(P);

%% bounds for BADS
%     lf        rf      ls      rs
pLB =[0.1,      0.5,    0.0001, 0.95];
pUB = [0.6,     0.8,    0.07,   0.9999];
LB = [0.0001,   0.5,    0.0001, 0.5];
UB = [0.9999,   0.9999, 0.15,   0.9999];
nruns = 10;

NS = [1:10,12:24];
for si = 1:length(NS)
    s = NS(si);
    file = sprintf('%s%sfits_%s_%02d.MAT',path,fits_folder,model,s);
    if exist(file,'file')~=2 || OVERWRITE
        load(sprintf('%s%ss%03d%s',path,data_folder,s,'.mat'))
        %% T structure for fitting
        T.trialtype = subject.TT;
        T.theta_planned = subject.TA;
        T.theta_planned(T.theta_planned == 0) = 360;
        T.Fext =subject.Fext;
        T.block =subject.block;
        T.nTrials = length(subject.TT);
        %% the update angle:
        %% Motion Based
        if strcmp(model,'motion_based')
            T.UpdateAngle =  -ThetaInDeg(s,:); % negative to map unit angle to robot angle
            %% Plan Based
        elseif strcmp(model,'plan_based')
            T.UpdateAngle = subject.TA;
        end
        %% Data & baseline correction of AI
        D = subject.r;
        for a =1:length(aPos)
            BL = mean(D(subject.TT == 5 &subject.TA == aPos(a)));
            ind = find(subject.TA == aPos(a) & logical(sum(subject.TT == TT.chan,2) ));
            D(ind) = D(ind) -BL;
        end
        % give force trials an AI of NaN
        D(logical(sum(subject.TT == TT.ff,2))) = NaN;
        %  exclude +- 60 degrees
        D(subject.TA == 60 | subject.TA == 300) = NaN;
        tic;
        X = NaN(nruns,length(pLB));
        fval = NaN(nruns,1);
        tim = NaN(nruns,1);
        for nrun = 1:nruns
            x0 = pLB+ rand(1,4).*(pUB-pLB);
            [X(nrun,:),fval(nrun)] = bads(@(x) GeneralizationFit(x,T,g,TT,D),x0,LB,UB,pLB,pUB);
            tim(nrun) = toc;
        end
        
        save(file,'X','fval','D','T')
        
    end
end