state_decoding_correct.m

%% decoder to decode state 1 vs state 2/3 from each task window

clear;
addpath(genpath('../multiarea_analysis'));
startup;
load('mycc.mat');
setup_colors;

datasheet = get_data_sheet('multiarea');

warning off; 
 
%% dataset information
opts = struct;
opts.base_dir = 'W:\Helmchen Group\Neurophysiology-Storage-03\Han\data\multiarea';
% opts.base_dir = '/home/ubuntu/neurophys/data/multiarea'; 
opts.data_dir = 'data_suite2p'; 
opts.result_dir = 'results_suite2p'; 
var_to_read = {'trial_vec', 'num_neuron', 'num_trial', 'trial_length', 'S_trial',  ...
    'ts', 'choice_time', 'task_label', 'F0', 'first_correct_lick'};

d1 = find_dataset(datasheet, 'experiment', '2AFC');
d2 = find_dataset(datasheet, 'experiment', '2AFC_no_cue');
dataset = cat(1, d1, d2);
dataset = dataset(dataset>=134);
dataset = dataset(dataset>=388 | dataset<=256);
dataset = sort(dataset, 'ascend');

x_str = {'texture', 'const'};
var_str = {'tex_vec', 'const_vec'};
global_model_result = 'behavior_glmhmm_all_logs_two_var'; 


result_name = 'state_decoding_correct_frame'; 

standard_trial_structure = [0.5,1,1,0.5,2,4];
standard_fr = 10;
standard_ts = make_trial_structure(standard_trial_structure, standard_fr);
standard_ts_fr = cellfun(@(x) x/standard_fr, standard_ts, 'uniformoutput', false);
trial_len = standard_ts{end}(end);
tvec = (1:trial_len)/standard_fr;
num_ts = 6;

tnum = 4;
tt_str = {'Correct 1', 'Correct 2', 'Incorrect 1', 'Incorrect 2'};

%% --- parameters for analysis ---
pstate_thr = 0;
% pstate_thr = 0.6;

method = 'frame';
% method = 'avg';

num_cv = 5;

% whether take pcs according to fix max number
pc_num_flag = 1;
num_pc_thr = 30;

% whether take pcs according to explained variance
pc_var_flag = 0;
expl_thr = 70;

remove_early_lick = true;

% whether to smooth data
smooth_flag = 1;
smooth_wsz = 3;



%% load model results
result_path = fullfile(opts.base_dir, 'results');
global_model = load(fullfile(result_path, [global_model_result '.mat']));

C = global_model.C;      % number of observation classes
D = global_model.D;      % number of GLM inputs (regressors)
K = global_model.K_best;      % number of latent states

local_file = sprintf('%s_learning_stage_A_only', global_model_result);

num_stage = 3;
stage_str = {'Naive', 'Learning', 'Expert'};

prediction_file = sprintf('%s_predict.mat', local_file);

%% train decoders
num_rep = 20;
num_shuff = 10;
nset = length(dataset);

for dataid = 1:nset

    fprintf('Processing dataset %d...\n', dataset(dataid));

    dinfo = data_info(datasheet, dataset(dataid), 'multiarea', opts.base_dir);
    spath = fullfile(dinfo.work_dir, opts.result_dir);
    if dinfo.quality_idx~=2; continue; end
    
    ld = load(fullfile(spath, prediction_file)); 
    p_state = ld.p_state; 
    state_pred = ld.state_pred; 
    keep_trials = ld.keep_idx; 

    state_pred = threshold_state_prob(state_pred, p_state, pstate_thr);
    
    num_state_trial = zeros(1, K);
    for k = 1:K
        num_state_trial(k) = sum(state_pred==k);
    end

    % combine state 2 and 3
    K_comb = 2;
    num_state_trial = [num_state_trial(1), sum(num_state_trial(2:3))];
    
    idx_state = cell(K_comb, 1);
    for k = 1:K_comb
        if k==1
            idx_state{k} = state_pred==1;
        else
            idx_state{k} = state_pred==2 | state_pred==3;
        end
    end

    %% load data
    data = load_data(dinfo, var_to_read, opts);
    data.trial_vec = data.trial_vec(keep_trials)';
    data.task_label.cue_vec = data.task_label.cue_vec(keep_trials);
    data.task_label.tex_vec = data.task_label.tex_vec(keep_trials);
    data.task_label.choice_vec = data.task_label.choice_vec(keep_trials);
    data.task_label.rew_vec = data.task_label.rew_vec(keep_trials);

    % correction for 187
    if isempty(data.ts{1}); data.ts{1} = 1:2; end
    data.ts{2} = data.ts{2}(data.ts{2}>0);
    data.first_correct_lick(data.first_correct_lick==0) = 1;

    % handle choice window
    data_out = process_choice_window_and_align_data(data, remove_early_lick);
    ts = data_out.ts;

    data_out.S_data = normalize_data_range(data_out.S_data);
    for a = 1:2
        data_out.S_data{a} = data_out.S_data{a}(:,keep_trials,:);
    end
    S_data = data_out.S_data;

    if smooth_flag
        for a = 1:2
            S_data{a} = smoothdata(S_data{a}, 3, 'gaussian', smooth_wsz);
        end
    end

    %% task labels
    label_str = {'tone', 'texture', 'choice', 'reward', 'ITI'};
    decoder_tw = {ts{2}, ts{3}, ts{4}, ts{5}, ts{6}};
    num_tw = length(decoder_tw);
    
   
    %% find trials in states  
   state_tt_idx = cell(K_comb, 1);
    for k = 1:K_comb
        t_idx = data.trial_vec==1 | data.trial_vec==2;
        state_tt_idx{k} = find(idx_state{k} & t_idx);
    end

    num_trial = cellfun(@(x) length(x), state_tt_idx);
    min_trial = min(num_trial(:));
    if min_trial<5; continue; end

    % bootstrap each trial type to match trial numbers
    t_idx_bs = cell(num_rep, K_comb);
    for rep_idx = 1:num_rep
        for k = 1:K_comb
            rand_idx = randperm(num_trial(k));
            rand_idx = rand_idx(1:min_trial);
            t_idx_bs{rep_idx,k} = state_tt_idx{k}(rand_idx);
        end
    end
    total_trial = length(t_idx_bs{1,1});

    % split trials into cross-validation sets
    idx_set = make_cv_set(min_trial, num_cv);

    % split train and test set indices
    num_cv = length(idx_set);
    idx_train = cell(num_cv, 1);
    idx_test = cell(num_cv, 1);
    for cv_idx = 1:num_cv
        idx_test{cv_idx} = idx_set{cv_idx};
        idx_train{cv_idx} = cell2mat(idx_set(setdiff(1:num_cv, cv_idx)));
    end

    %% ensure enough trials before doing the heavy computations
    if total_trial<5; continue; end

    %% do pca on dataset
    explained_var = zeros(1, 2);
    num_pc = zeros(1, 2);

    S_pc = cell(1,2);
    if pc_num_flag
        for a = 1:2
            [S_pc(a), explained_var(a), num_pc(a)] = pca_reconstruct...
                (S_data(a), 'num_pc', num_pc_thr, 'zscore', false);
        end

    elseif pc_var_flag
        for a = 1:2
            [S_pc(a), explained_var(a), num_pc(a)] = pca_reconstruct...
                (S_data(a), 'explained', expl_thr, 'zscore', false);
        end
    end

    
    %% train decoders for each state, shuffle outside of CVs
    Lambda = logspace(-5,1,7);
    fitbias = true;
    reg_str = 'ridge';
    
    task_decoder_auc = nan(num_tw, 2, num_rep);
    task_decoder_auc_shuff = nan(num_tw, 2, num_rep);
    task_decoder_b = cell(num_tw, 2); 

    for n = 1:num_tw
        tw = decoder_tw{n};
        for a = 1:2
            b_rep = nan(num_pc(a)+1, num_rep, num_cv);
            parfor rep_idx = 1:num_rep
%             for rep_idx = 1:num_rep

                % train on combined states
                trial_idx = cat(1, t_idx_bs{rep_idx,1}, t_idx_bs{rep_idx,2});

                X0 = S_pc{a}(:,trial_idx,tw);
                Y0 = cat(1, ones(min_trial,1), 2*ones(min_trial,1))';
                X0(isnan(X0)) = 0;
                
                if isempty(X0) || sum(Y0==1)==0 || sum(Y0==2)==0
                    continue; 
                end
                
                if strcmp(method, 'avg')
                    X = nanmean(X0, 3)';
                    Y = Y0;
                elseif strcmp(method, 'frame')
                    X = flatten_trial_trace(X0)';
                    Y = reshape(repmat(Y0, length(tw), 1), [], 1);
                end

                % find best lambda
                C = [0,length(Y)/sum(Y==1); length(Y)/sum(Y==2),0];
                mdl = fitclinear(X, Y, 'ClassNames',[1,2], 'Lambda', Lambda, ...
                    'Prior','Uniform', 'Cost', C, 'regularization', reg_str, ...
                    'Kfold', num_cv, 'FitBias', fitbias);
                ce = kfoldLoss(mdl); 
                [~,idx] = min(ce);
                lambda_idx = Lambda(idx);
                
                % train k-fold on real data
                auc_cv = nan(1, num_cv);
                auc_s = nan(num_shuff, num_cv);
                for cv_idx = 1:num_cv
                    i_train = cat(1, idx_train{cv_idx}, idx_train{cv_idx}+total_trial);
                    i_test = cat(1, idx_test{cv_idx}, idx_test{cv_idx}+total_trial);

                    if strcmp(method, 'avg')
                        x_train = X(i_train,:);
                        y_train = Y(i_train); 
                        x_test = X(i_test,:);
                        y_test = Y(i_test); 
                    elseif strcmp(method, 'frame')
                        x_train = flatten_trial_trace(X0(:,i_train,:))';
                        y_train = reshape(repmat(Y0(i_train), length(tw),1), [], 1);
                        x_test = flatten_trial_trace(X0(:,i_test,:))';
                        y_test = reshape(repmat(Y0(i_test), length(tw),1), [], 1);
                    end
                    if sum(y_train==1)==0 || sum(y_train==2)==0
                        continue;
                    end
                    C = [0,length(y_train)/sum(y_train==1); length(y_train)/sum(y_train==2),0];
                    b = fitclinear(x_train, y_train, 'ClassNames',[1,2], 'Lambda',lambda_idx,...
                        'Prior', 'Uniform', 'Cost', C, 'regularization', reg_str, 'FitBias', fitbias);
                    % task_decoder_b{n,k,a,rep_idx,cv_idx} = [b.Bias; b.Beta];
                    b_rep(:,rep_idx,cv_idx) = [b.Bias; b.Beta];
                    
                    yhat = x_test*b.Beta;
                    auc_cv(cv_idx) = scoreAUC(y_test(:)==2, yhat);

                    % apply model to shuffled data
                    for s = 1:num_shuff
                        rand_idx = randperm(size(X,2));  % shuffled neuron
                        if strcmp(method, 'avg')
                            xs_test = X(i_test,rand_idx);
                        elseif strcmp(method, 'frame')
                            xs_test = flatten_trial_trace(X0(rand_idx,i_test,:))';
                        end
                        if sum(y_test==1)==0 || sum(y_test==2)==0
                            continue;
                        end
                        % auc             
                        yhat = xs_test*b.Beta;
                        auc_s(s,cv_idx) = scoreAUC(y_test(:)==2, yhat);
                    end 

                end
                auc_s = nanmean(auc_s, 2);
                task_decoder_auc_shuff(n,a,rep_idx) = quantile(auc_s, 0.95);
                task_decoder_auc(n,a,rep_idx) = nanmean(auc_cv);
            end
            task_decoder_b{n,a} = b_rep;
        end
    end
    
    
    %% save
    result_file = fullfile(spath, sprintf('%s.mat', result_name));
    save(result_file, 'task_decoder_auc', 'task_decoder_auc_shuff', ...
        'task_decoder_b', 'min_trial', 't_idx_bs', 't_idx', 'idx_set', ...
        'num_cv',  'smooth_flag', 'smooth_wsz', ...
        'method', 'pstate_thr', ...
        'pc_num_flag', 'num_pc_thr', 'pc_var_flag', 'expl_thr', '-v7.3');


    %% some plotting 
    plot_result = 0;
    if plot_result
    
    %% plot predicted auc
    cc = {mycc.orange, mycc.green};
    figure;
    for a = 1:2
        subplot(1, 2, a); hold on;
        vs = squeeze(task_decoder_auc_shuff(:,a,:));
        ym = nanmean(vs, 2); yse = nanstd(vs,[],2)/sqrt(num_rep);
        confplot(1:num_tw, ym', yse', yse', mycc.gray, 0.2);

        v = squeeze(task_decoder_auc(:,a,:));
        ym = nanmean(v, 2); yse = nanstd(v,[],2)/sqrt(num_rep);
        confplot(1:num_tw, ym', yse', yse', mycc.black, 0.2);
        
        set(gca, 'xtick', 1:num_tw, 'xticklabel', label_str);
        title(sprintf('A%d', a));
    end
    linkaxes;
    
    end
    
end