clc, clear
clear global
rng(0,'twister');
global soft_out
options = optimset;
options = optimset(options,'Display', 'off','PlotFcns', {@optimplotx @optimplotfval}, ...
   'Algorithm', 'active-set', 'FinDiffType', 'central');
options.PlotFcns = [];
D=dir('*.csv');
subjStay=[];
subjProbe=[];
subjHabit=[];
subjLearning=[];
AIC=[];
for s=1:length({D.name})
    %% Data organizing
    format longg
    opts = detectImportOptions(D(s).name);
    opts = setvartype(opts, 'condition', 'char');
    opts = setvartype(opts, 'phase', 'char');
    opts = setvartype(opts, 'block_type', 'char');
    opts = setvartype(opts, 'left_shape_reward_level', 'char');
    opts = setvartype(opts, 'right_shape_reward_level', 'char');
    dat = readtable(D(s).name, opts);
    
    % Shape assignments
    shapesCB = erase(split(dat.shapes(1))', ["[", "]", "'", ","]);
    shapesSocial = erase(split(dat.shapesSocial(1))', ["[", "]", "'", ","]);
    shapesIndiv = erase(split(dat.shapesIndividual(1))', ["[", "]", "'", ","]);
    
    % Core variables
    choice = dat.choice(~isnan(dat.choice));
    trialType = dat.trial_type(~isnan(dat.choice));% 1=normal, 0=probe
    rewLev = dat.rewardLevel(~isnan(dat.choice));
    origRewLev = dat.original_reward_level(~isnan(dat.choice));
    phase = dat.phase(~isnan(dat.choice));
    blockType = dat.block_type(~isnan(dat.choice));
    leftShapeRew = dat.left_shape_reward_level(~isnan(dat.choice));
    rightShapeRew = dat.right_shape_reward_level(~isnan(dat.choice));
    rewLeft = dat.rewardLeft(~isnan(dat.choice));
    rewRight = dat.rewardRight(~isnan(dat.choice));
    imgLeft = dat.imageLeft(~isnan(dat.choice));
    imgRight = dat.imageRight(~isnan(dat.choice));
    cons = dat.payoffCons(~isnan(dat.choice));% consensus level (0,1,2)
    
    chosUnchosOPT = zeros(height(choice),3);% chosen, unchosen, block_type(1=social,2=indiv)
    chosUnchosREW = [rewLeft, rewRight];
    for cb = 1:8
        chosUnchosOPT(strcmp(imgLeft, shapesCB{cb}),1) = cb;
        chosUnchosOPT(strcmp(imgRight, shapesCB{cb}),2) = cb;
    end
    chosUnchosOPT(strcmp(blockType, 'social'),3) = 1;
    chosUnchosOPT(strcmp(blockType, 'individual'),3) = 2;
    
    % Convert to chosen/unchosen
    for i = 1:height(choice)
        if choice(i)==39
            chosUnchosOPT(i,1:2) = [chosUnchosOPT(i,2), chosUnchosOPT(i,1)];
            chosUnchosREW(i,1:2) = [chosUnchosREW(i,2), chosUnchosREW(i,1)];
        end
    end
    
    %% Behavioral Analyses
    % Probe choices - proportion choosing high reward shape
    % Indices: phase(1=train,2=test) x blockType(1=social,2=indiv)
    probeHigh = zeros(2,2);
    probeTotal = zeros(2,2);
    for i = 1:height(choice)
        if trialType(i)==0% probe trial
            phaseIdx = 1 + strcmp(phase{i}, 'test');
            blockIdx = chosUnchosOPT(i,3);
            probeTotal(phaseIdx, blockIdx) = probeTotal(phaseIdx, blockIdx) + 1;
            if (choice(i)==37 && strcmp(leftShapeRew{i}, 'high')) || ...
               (choice(i)==39 && strcmp(rightShapeRew{i}, 'high'))
                probeHigh(phaseIdx, blockIdx) = probeHigh(phaseIdx, blockIdx) + 1;
            end
        end
    end
    probeProp = probeHigh ./ probeTotal;% [train;test] x [social,indiv]
    
    % Stay probabilities
    % Indices: rewLev(1=low,2=high) x blockType(1=social,2=indiv) x phase(1=train,2=test)
    bigStay = zeros(2,2,2);
    numTrials = zeros(2,2,2);
    % Win/lose stay: outcome(1=win,2=lose) x blockType x phase
    winStay = zeros(2,2,2);
    winTrials = zeros(2,2,2);
    
    for i = 1:height(trialType)-1
        if trialType(i)==1% normal/feedback trial
            if trialType(i+1)==1% next is also normal
                % Check if chosen shape available on next trial
                chosenShape = chosUnchosOPT(i,1);
                if chosUnchosOPT(i+1,1)==chosenShape || chosUnchosOPT(i+1,2)==chosenShape
                    phaseIdx = 1 + strcmp(phase{i}, 'test');
                    blockIdx = chosUnchosOPT(i,3);
                    rewIdx = rewLev(i) + 1;
                    winIdx = 1 + (chosUnchosREW(i,1) < chosUnchosREW(i,2));% 1=win, 2=lose
                    
                    numTrials(rewIdx, blockIdx, phaseIdx) = numTrials(rewIdx, blockIdx, phaseIdx) + 1;
                    winTrials(winIdx, blockIdx, phaseIdx) = winTrials(winIdx, blockIdx, phaseIdx) + 1;
                    
                    if chosUnchosOPT(i+1,1)==chosenShape% stayed
                        bigStay(rewIdx, blockIdx, phaseIdx) = bigStay(rewIdx, blockIdx, phaseIdx) + 1;
                        winStay(winIdx, blockIdx, phaseIdx) = winStay(winIdx, blockIdx, phaseIdx) + 1;
                    end
                end
            end
        end
    end
    behavStay = bigStay ./ numTrials;
    behavStayWin = winStay ./ winTrials;
    
    % Learning curves - test phase normal trials, binned
    testNormIdx = find(trialType==1 & strcmp(phase, 'test'));
    testNormSocial = find(trialType==1 & strcmp(phase, 'test') & strcmp(blockType, 'social'));
    testNormIndiv = find(trialType==1 & strcmp(phase, 'test') & strcmp(blockType, 'individual'));
    
    % Track choosing currently-optimal (higher reward)
    chooseOpt = zeros(length(testNormIdx),1);
    for j = 1:length(testNormIdx)
        i = testNormIdx(j);
        if choice(i)==37
            chooseOpt(j) = rewLeft(i) >= rewRight(i);
        else
            chooseOpt(j) = rewRight(i) >= rewLeft(i);
        end
    end
    
    % Bin into thirds
    nTest = length(testNormIdx);
    binSize = floor(nTest/3);
    if nTest >= 3
        learnCurve = [mean(chooseOpt(1:binSize)), ...
                      mean(chooseOpt(binSize+1:2*binSize)), ...
                      mean(chooseOpt(2*binSize+1:end))];
    else
        learnCurve = [NaN, NaN, NaN];
    end
    
    % Learning curves by block type
    chooseOptSocial = zeros(length(testNormSocial),1);
    chooseOptIndiv = zeros(length(testNormIndiv),1);
    for j = 1:length(testNormSocial)
        i = testNormSocial(j);
        chooseOptSocial(j) = (choice(i)==37 && rewLeft(i)>=rewRight(i)) || ...
                             (choice(i)==39 && rewRight(i)>=rewLeft(i));
    end
    for j = 1:length(testNormIndiv)
        i = testNormIndiv(j);
        chooseOptIndiv(j) = (choice(i)==37 && rewLeft(i)>=rewRight(i)) || ...
                            (choice(i)==39 && rewRight(i)>=rewLeft(i));
    end
    
    nSoc = length(testNormSocial); nInd = length(testNormIndiv);
    binSoc = floor(nSoc/3); binInd = floor(nInd/3);
    if nSoc >= 3
        learnSocial = [mean(chooseOptSocial(1:binSoc)), ...
                       mean(chooseOptSocial(binSoc+1:2*binSoc)), ...
                       mean(chooseOptSocial(2*binSoc+1:end))];
    else
        learnSocial = [NaN, NaN, NaN];
    end
    if nInd >= 3
        learnIndiv = [mean(chooseOptIndiv(1:binInd)), ...
                      mean(chooseOptIndiv(binInd+1:2*binInd)), ...
                      mean(chooseOptIndiv(2*binInd+1:end))];
    else
        learnIndiv = [NaN, NaN, NaN];
    end
    
    % Store subject specs
    subjSpecs{s} = {trialType, chosUnchosOPT, chosUnchosREW, phase, blockType, rewLev, origRewLev, cons};
    
    % Aggregate subject data
    % probeProp: [train;test] x [social,indiv]
    subjProbe(s,:) = [probeProp(1,1), probeProp(1,2), probeProp(2,1), probeProp(2,2)];% train_soc, train_ind, test_soc, test_ind
    
    % behavStay: rewLev x blockType x phase -> flatten to [low_soc_train, low_ind_train, high_soc_train, high_ind_train, low_soc_test, ...]
    subjStay(s,:) = [squeeze(behavStay(1,:,1)), squeeze(behavStay(2,:,1)), ...
                     squeeze(behavStay(1,:,2)), squeeze(behavStay(2,:,2))];
    
    % behavStayWin: outcome x blockType x phase
    subjStayWin(s,:) = [squeeze(behavStayWin(1,:,1)), squeeze(behavStayWin(2,:,1)), ...
                        squeeze(behavStayWin(1,:,2)), squeeze(behavStayWin(2,:,2))];% win_soc_train, win_ind_train, lose_soc_train, ...
    
    % Habit index: test probe high-choice (higher = more habitual)
    subjHabit(s,:) = [probeProp(2,1), probeProp(2,2), ...% raw habit
                      probeProp(2,1)-0.5, probeProp(2,2)-0.5, ...% habit strength
                      probeProp(1,1)-0.5, probeProp(1,2)-0.5];% learning strength
    
    % Learning curves
    subjLearning(s,:) = learnCurve;
    subjLearnSocial(s,:) = learnSocial;
    subjLearnIndiv(s,:) = learnIndiv;
    
%% Fit & Compare Models
    k=[4,3,3,2];% # of parameters for AIC: hybrid, imitation, social, baseline
    for mods=1:4
        paramRec=0;
        numPar=repmat([0.5; 0.01; 0.99], 1, k(mods)-1);
        [params, nll, ~, ~]=fminsearchbnd(@(fp) fitparams(fp, subjSpecs{s}, mods, paramRec),[numPar(1,:),5],[numPar(2,:),1],[numPar(3,:),10],options);
        subjParams{s,mods}=params;
        AIC(s,mods)=2*k(mods)+2*(nll);
    end
end
save habitSpecs subj* AIC

%% Group Results
fprintf('\n=== PROBE CHOICES (P choose high) ===\n');
fprintf('Training - Social: %.3f (%.3f), Individual: %.3f (%.3f)\n', ...
    mean(subjProbe(:,1)), std(subjProbe(:,1)), mean(subjProbe(:,2)), std(subjProbe(:,2)));
fprintf('Test - Social: %.3f (%.3f), Individual: %.3f (%.3f)\n', ...
    mean(subjProbe(:,3)), std(subjProbe(:,3)), mean(subjProbe(:,4)), std(subjProbe(:,4)));

fprintf('\n=== HABIT STRENGTH (vs chance) ===\n');
fprintf('Social: %.3f (%.3f), Individual: %.3f (%.3f)\n', ...
    mean(subjHabit(:,3)), std(subjHabit(:,3)), mean(subjHabit(:,4)), std(subjHabit(:,4)));

fprintf('\n=== LEARNING CURVES (test phase) ===\n');
fprintf('Social - Early: %.3f, Mid: %.3f, Late: %.3f\n', ...
    nanmean(subjLearnSocial(:,1)), nanmean(subjLearnSocial(:,2)), nanmean(subjLearnSocial(:,3)));
fprintf('Individual - Early: %.3f, Mid: %.3f, Late: %.3f\n', ...
    nanmean(subjLearnIndiv(:,1)), nanmean(subjLearnIndiv(:,2)), nanmean(subjLearnIndiv(:,3)));

fprintf('\n=== MODEL COMPARISON (AIC) ===\n');
fprintf('Hybrid: %.2f, Imitation: %.2f, Social: %.2f, Baseline: %.2f\n', ...
    mean(AIC(:,1)), mean(AIC(:,2)), mean(AIC(:,3)), mean(AIC(:,4)));

%% Models
function nll=fitparams(fp, specs, mods, paramRec)
    global soft_out
    % specs: 1.trialType, 2.chosUnchosOPT, 3.chosUnchosREW, 4.phase, 5.blockType, 6.rewLev, 7.origRewLev, 8.cons
    % chosUnchosOPT: col1=chosen shape, col2=unchosen shape, col3=blockType(1=social,2=indiv)
    
    freqVals = ones(1,8);% frequency bonus for each shape
    if mods<3 % hybrid and imitation
        actVals = zeros(1,8);% action imitation values
        payVals = zeros(1,8)+0.5;% payoff values
        sumVals = actVals+payVals;
    else
        sumVals = zeros(1,8)+0.5;
    end
    
    nTrials = height(specs{1});
    M = zeros(nTrials,2);
    modelprob = zeros(nTrials,1);
    specsO = specs;
    
    for i = 1:nTrials
        opt1 = specs{2}(i,1);% chosen option
        opt2 = specs{2}(i,2);% unchosen option
        
        M(i,1) = sumVals(opt1) + sqrt(2*log(i)/freqVals(opt1));
        M(i,2) = sumVals(opt2) + sqrt(2*log(i)/freqVals(opt2));
        modelprob(i) = exp(fp(end)*M(i,1)) / (exp(fp(end)*M(i,1)) + exp(fp(end)*M(i,2)));
        
        if paramRec==1
            chosResp = randsample([1 0], 1, true, [modelprob(i) 1-modelprob(i)]);
            if chosResp==0
                specs{2}(i,1:2) = [specsO{2}(i,2), specsO{2}(i,1)];
                specs{3}(i,1:2) = [specsO{3}(i,2), specsO{3}(i,1)];
                specs{8}(i) = 2 - specsO{8}(i);% flip consensus
            end
        end
        
        if specs{1}(i)==1% normal/feedback trial - update values
            freqVals(specs{2}(i,1)) = freqVals(specs{2}(i,1)) + 1;
            cons = specs{8}(i);% actual consensus from data (0,1,2)
            
            if mods==1% Hybrid model
                actVals(specs{2}(i,1)) = actVals(specs{2}(i,1)) + fp(2)*(cons - actVals(specs{2}(i,1)));
                payVals(specs{2}(i,1)) = payVals(specs{2}(i,1)) + fp(1)*(specs{3}(i,1) - payVals(specs{2}(i,1))) + fp(3)*(cons - payVals(specs{2}(i,1)));
                sumVals = actVals + payVals;
            elseif mods==2% Imitation model
                actVals(specs{2}(i,1)) = actVals(specs{2}(i,1)) + fp(2)*(cons - actVals(specs{2}(i,1)));
                payVals(specs{2}(i,1)) = payVals(specs{2}(i,1)) + fp(1)*(specs{3}(i,1) - payVals(specs{2}(i,1)));
                sumVals = actVals + payVals;
            elseif mods==3% Social Reward model
                sumVals(specs{2}(i,1)) = sumVals(specs{2}(i,1)) + fp(1)*(specs{3}(i,1) - sumVals(specs{2}(i,1))) + fp(2)*(cons - sumVals(specs{2}(i,1)));
            elseif mods==4% Money Only model
                sumVals(specs{2}(i,1)) = sumVals(specs{2}(i,1)) + fp(1)*(specs{3}(i,1) - sumVals(specs{2}(i,1)));
            end
        end
    end
    
    nll = sum(log(modelprob))*-1;
    if paramRec == 1
        soft_out = specs;
    end
end
