state_task_decoding.m

%% decoder to decode task variables 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.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_task_decoding_combined_correct_pc30_frame'; 


%% --- parameters for analysis ---
param = struct();

param.pstate_thr = 0;
% param.pstate_thr = 0.6;

param.method = 'frame';
% param.method = 'avg';

param.num_rep = 20;
param.num_cv = 5;
param.num_shuff = 100;

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

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

param.remove_early_lick = true;

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

% extract some parameters for convenience
num_rep = param.num_rep;
num_cv = param.num_cv;
num_shuff = param.num_shuff;

%% 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
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, param.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, param.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 param.smooth_flag
        for a = 1:2
            S_data{a} = smoothdata(S_data{a}, 3, 'gaussian', param.smooth_wsz);
        end
    end

    %% task labels
    label = {data.task_label.cue_vec, ...
            data.task_label.tex_vec, ...
            data.task_label.choice_vec};
    label_str = {'tone', 'texture', 'choice'};
    decoder_tw = {ts{2}, ts{3}, ts{4}};

    num_decoder = length(label);
    
   
    %% 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 and stratify to match conditions
    % (first half is state 1, second half is state 2 in t_idx_bs)
    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 param.pc_num_flag
        for a = 1:2
            [S_pc(a), explained_var(a), num_pc(a)] = pca_reconstruct...
                (S_data(a), 'num_pc', param.num_pc_thr, 'zscore', false);
        end

    elseif param.pc_var_flag
        for a = 1:2
            [S_pc(a), explained_var(a), num_pc(a)] = pca_reconstruct...
                (S_data(a), 'explained', param.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';
    method = param.method;
    
    task_decoder_auc = nan(num_decoder, K_comb, 2, num_rep);
    task_decoder_auc_shuff = nan(num_decoder, K_comb, 2, num_rep);
    task_decoder_b = cell(num_decoder, 2); 

    for n = 1:num_decoder
        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 = label{n}(trial_idx);
                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(K_comb, num_cv);
                auc_s = nan(K_comb, num_shuff, num_cv);
                for cv_idx = 1:num_cv
                    idx = cat(1, idx_train{cv_idx}, idx_train{cv_idx}+total_trial);
                    if strcmp(method, 'avg')
                        x_train = X(idx,:);
                        y_train = Y(idx); 
                    elseif strcmp(method, 'frame')
                        x_train = flatten_trial_trace(X0(:,idx,:))';
                        y_train = reshape(repmat(Y0(idx), length(tw),1), [], 1);
                    end
                    if sum(y_train==1)==0 || sum(y_train==2)==0
%                         b_rep(:,rep_idx,cv_idx) = nan(data.num_neuron(a)+1,1);
                        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];
                    
                    % auc
                    for k = 1:K_comb
                        idx = idx_test{cv_idx} + (k-1)*total_trial;
                        if strcmp(method, 'avg')
                            x_test = X(idx,:);
                            y_test = Y(idx); 
                        elseif strcmp(method, 'frame')
                            x_test = flatten_trial_trace(X0(:,idx,:))';
                            y_test = reshape(repmat(Y0(idx), length(tw),1), [], 1);
                        end
                        if sum(y_test==1)==0 || sum(y_test==2)==0; continue; end
                        yhat = x_test*b.Beta;
                        auc_cv(k,cv_idx) = scoreAUC(y_test==2, yhat);
                    end
                    

                    % train shuffled data
                    for s = 1:num_shuff
                        Ys = Y(randperm(length(Y)));  % shuffled label
                        idx = cat(1, idx_train{cv_idx}, idx_train{cv_idx}+total_trial);
                        if strcmp(method, 'avg')
                            xs_train = X(idx,:);
                            ys_train = Ys(idx); 
                        elseif strcmp(method, 'frame')
                            xs_train = flatten_trial_trace(X0(:,idx,:))';
                            ys_train = reshape(repmat(Ys(idx), length(tw),1), [], 1);
                            ys_train = ys_train(randperm(length(ys_train)));
                        end
                        if sum(ys_train==1)==0 || sum(ys_train==2)==0
                            continue;
                        end
                        C = [0,length(ys_train)/sum(ys_train==1); ...
                            length(ys_train)/sum(ys_train==2),0];
                        bs = fitclinear(xs_train, ys_train, 'ClassNames',[1,2], 'Lambda',lambda_idx,...
                            'Prior', 'Uniform', 'Cost', C, 'regularization', reg_str, 'FitBias', fitbias);
                        
                        % auc
                        for k = 1:K_comb
                            idx = idx_test{cv_idx} + (k-1)*total_trial;
                            if strcmp(method, 'avg')
                                xs_test = X(idx,:);
                                ys_test = Ys(idx); 
                            elseif strcmp(method, 'frame')
                                xs_test = flatten_trial_trace(X0(:,idx,:))';
                                ys_test = reshape(repmat(Ys(idx), length(tw),1), [], 1);
                                ys_test = ys_test(randperm(length(ys_test)));
                            end
                            if sum(ys_test==1)==0 || sum(ys_test==2)==0; continue; end
                            yhat = xs_test*bs.Beta;
                            auc_s(k,s,cv_idx) = scoreAUC(ys_test==2, yhat);
                        end
                    end 
                end
                auc_s = nanmean(auc_s, 3);
                task_decoder_auc_shuff(n,:,a,rep_idx) = quantile(auc_s, 0.95, 2);
                task_decoder_auc(n,:,a,rep_idx) = nanmean(auc_cv, 2);
            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', ...
        'param', '-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;
        for k = 1:K_comb
            vs = squeeze(task_decoder_auc_shuff(:,k,a,:));
            ym = nanmean(vs, 2); yse = nanstd(vs,[],2)/sqrt(num_rep);
            confplot(1:num_decoder, ym', yse', yse', lighten_color(cc{k}, 0.5), 0.2);

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