%% Load data
% 
% wh_comp = 'byeolmac';
% [~, ~, ~, fastprojectdir] = fast_vary_path(wh_comp);
% [D, fast_sub, jth, sub_f] = fast_load_data(fastprojectdir, 'behav-sess1');
% 
% figdir = fullfile(fastprojectdir, 'figures','Markov chain');
% tabledir = fullfile(fastprojectdir, 'figures','table');
% datdir = fullfile(fastprojectdir, 'data');
% 
% [~, ~, ~, web_syncdir] = fastweb_vary_path(wh_comp);
% [~, fastweb_sub, ~, ~] = fastweb_load_data(web_syncdir);
% web = load(fullfile(web_syncdir, 'data','fastweb_210125_markov_final_model.mat'));   % valence, self, time
% 
% trainid_web = find(fastweb_sub.D_Happiness_n137)'; 
% testid_web = fastweb_sub.twice_first_n186';     
% 
% subj_j = [fast_sub.beh.firstsess_n93; fast_sub.beh.secondsess_n93];
% [s, idx] = sort(fast_sub.beh.twice_second_n93);
% first_idx = fast_sub.beh.twice_first_n93(idx);
% for i = 1:numel(first_idx)
%     first_idx_n92(i,1) = find(subj_j == first_idx(i));
% end
% testid = 63:92; % second session
% 
% dimensions = {'valence', 'safe_threat', 'time',  'self-related', 'vividness'};
% dat = get_var(D, dimensions, 'conditional', {'eventname', 'view'});
% questionnaire_names = D.Subj_Level.names([19:20 18 23:25 21 27 29]);
% q_dat = get_var(D, questionnaire_names);
% q_dat = q_dat(subj_j,:);

datadir = '/Users/byeolkim/dropbox/FAST/sync/data/public sharing';
savename = fullfile(datadir, 'figure2_a.mat');
load(savename);

%% Create features for training models from Markov chain

% 1. mean(scores): mean scores 1 x 5 
% 2. var(scores): how much vary 1 x 5 
% 3. transition rate 9 x 3, 4 x 2
% 4. steady state prob 3 x 3, 2 x 2

dimensions = {'valence', 'safe_threat', 'time',  'self-related', 'vividness'};
markov = cell(1,numel(dimensions));

thresh.neg = -.33;
thresh.pos = .33;
thresh.half = .5;

for dim_i = 1:numel(dimensions)
    j = 0;
    for subj_i = subj_j' % subj_i = 1
        j = j + 1;
        
        idx = zeros(160,1);
        dati_val = dat(subj_i,:,dim_i)';
        
        features_raw{dim_i}{j} = var(dati_val);                             % 1. variance
        features_raw{dim_i}{j} = [features_raw{dim_i}{j}; mean(dati_val)];      % 2. mean
        
        if dim_i < 4 % valence, safe_threat, time
            idx(dati_val< thresh.neg)=1;
            idx(dati_val>= thresh.neg & dati_val< thresh.pos)=2;
            idx(dati_val>= thresh.pos)=3;
            n_state = 3;
        else
            idx(dati_val< thresh.half) = 1;
            idx(dati_val>= thresh.half) = 2;
            n_state = 2;
        end
        
        out = markov_chain(idx, 'n_state', n_state,'boundary', 40:40:120);
        
        trans_prob = out{1}.trans_prob(:);
        features_raw{dim_i}{j} = [features_raw{dim_i}{j}; trans_prob]; % 3. trans_prob
        
        steady_state_prob = out{1}.steady_state_prob';
        features_raw{dim_i}{j} = [features_raw{dim_i}{j}; steady_state_prob]; % 4. steady_state_prob
        
        if j == 1
            features_descrip{dim_i} = repmat({['variance_' dimensions{dim_i}]}, size(var(dati_val)));
            features_descrip{dim_i} = [features_descrip{dim_i}; repmat({['mean_' dimensions{dim_i}]}, size(mean(dati_val)))];
            features_descrip{dim_i} = [features_descrip{dim_i}; repmat({['trans_prob_' dimensions{dim_i}]}, size(trans_prob))];
            features_descrip{dim_i} = [features_descrip{dim_i}; repmat({['steady_state_prob_' dimensions{dim_i}]}, size(steady_state_prob))];
        end
    end
end

% more annotation:
trans_prob_dir_9 = {'11', '21', '31', '12', '22', '32', '13', '23', '33'}; trans_prob_dir_4 = {'11', '21', '12', '22'};
state_state_3 = {'1', '2', '3'}; state_state_2 = {'1', '2'};

for dim_i = 1:numel(dimensions)
    if  dim_i <=3
        for i = 3:11
            features_descrip{dim_i}{i} = [features_descrip{dim_i}{i} '_' trans_prob_dir_9{i-2}];
        end
    else
        for i = 3:6
            features_descrip{dim_i}{i} = [features_descrip{dim_i}{i} '_' trans_prob_dir_4{i-2}];
        end
    end
    
    if  dim_i <=3
        for i = 12:14
            features_descrip{dim_i}{i} = [features_descrip{dim_i}{i} '_' state_state_3{i-11}];
        end
    else
        for i = 7:8
            features_descrip{dim_i}{i} = [features_descrip{dim_i}{i} '_' state_state_2{i-6}];
        end
    end    
end


%% Features of FIRST SESSION PARTICIPANTS
trainid = 1:62; 

% x = stack all features across the first session subjects [58 x 62]
k = 0;
for j = trainid
    k = k + 1;    
    feature_temp = [];
    for dim_i = 1:numel(dimensions)
        a = features_raw{dim_i}{j}(:);
        feature_temp = [feature_temp; a(:)]; 
    end
    feature_trained(:,k) = feature_temp;     
end

features_descrip_all = [features_descrip{1}(:); features_descrip{2}(:); features_descrip{3}(:); features_descrip{4}(:); features_descrip{5}(:)];

%% Extract Factors in questionnaire responses (Table_S2)
q_negativity.sess1 = q_dat(trainid,:);

% factor analysis with personality of the first session participants
Tgt = [1 0 0 0 1 0 0 0 0; 0 1 1 1 0 1 1 1 1]'; 

[lambda,psi,T,~,F_neg.sess1] = factoran(scale(q_negativity.sess1), 2, 'rotate', 'pattern', 'target', Tgt);
table(questionnaire_names', lambda)  % Table_S2
F_descript = {'negativity','positivity'};

%                  Var1                         lambda        
%     ______________________________    ______________________ 
%     {'PANAS_PA'                  }     0.080843      0.93911
%     {'PANAS_NA'                  }      0.73607      0.26444
%     {'CES-D'                     }      0.66709     -0.32901
%     {'MASQ30_general'            }      0.86947    -0.072784
%     {'MASQ30_anhedonicdepression'}    -0.093794      0.80945
%     {'MASQ30_anxietyarousal'     }      0.57112      0.32308
%     {'STAI-T'                    }      0.79391     -0.24883
%     {'RRS_brooding'              }       0.7471     0.059702
%     {'RRS_depressiverumination'  }      0.81406     0.069459

%% 
% >>>>>>>>>>>>>>>> TABLE >>>>>>>>>>>>>>>>>>>>
% 
%% Table S1: Stability of the features across different sets of seed words and test time points (7-week interval)
% reliability between features of the first session and second session

ft_sess1 = feature_trained(:,first_idx_n92);       % features of the first session participants

clear x2;
k = 0;
for j = testid
    k = k + 1;
    feature_temp = [];
    for dim_i = 1:numel(dimensions)
        a = features_raw{dim_i}{j}';
        feature_temp = [feature_temp; a(:)]; 
    end
    ft_sess2(:,k) = feature_temp;       
end

for i = 1:size(ft_sess1,1)
    [test_retest.reliability_x(i,1), test_retest.reliability_x_p(i,1)] = corr(ft_sess1(i,:)', ft_sess2(i,:)');
end
asterisk = pnum2star(test_retest.reliability_x_p);
ind = [2 1 3:14 16 15 17:28 30 29 31:42 44 43 45:50 52 51 53:58];
 
t = table(features_descrip_all(ind), test_retest.reliability_x(ind), ...
    asterisk(ind), test_retest.reliability_x_p(ind), 'VariableNames',{'features','r','asterisk','p'})

%   58×4 table
%                   features                      r          asterisk          p     
%     ____________________________________    _________    ____________    __________
% 
%     {'mean_valence'                    }      0.54304    {'**'      }     0.0019297
%     {'variance_valence'                }      0.72112    {'***'     }    6.9374e-06
%     {'trans_prob_valence_11'           }       0.4076    {'*'       }      0.025366
%     {'trans_prob_valence_21'           }      0.45611    {'*'       }      0.011301
%     {'trans_prob_valence_31'           }      0.58948    {'***'     }    0.00060845
%     {'trans_prob_valence_12'           }      0.56722    {'**'      }     0.0010805
%     {'trans_prob_valence_22'           }      0.57864    {'***'     }    0.00080902
%     {'trans_prob_valence_32'           }      0.48878    {'**'      }     0.0061301
%     {'trans_prob_valence_13'           }      0.50486    {'**'      }     0.0044363
%     {'trans_prob_valence_23'           }      0.43816    {'*'       }      0.015445
%     {'trans_prob_valence_33'           }      0.55385    {'**'      }     0.0014969
%     {'steady_state_prob_valence_1'     }      0.68994    {'***'     }    2.4626e-05
%     {'steady_state_prob_valence_2'     }      0.67338    {'***'     }    4.5405e-05
%     {'steady_state_prob_valence_3'     }      0.56598    {'**'      }     0.0011143
%     {'mean_safe_threat'                }      0.77214    {'***'     }    5.8043e-07
%     {'variance_safe_threat'            }      0.75426    {'***'     }    1.4812e-06
%     {'trans_prob_safe_threat_11'       }      0.27149    {0×0 double}       0.14671
%     {'trans_prob_safe_threat_21'       }      0.74117    {'***'     }    2.8009e-06
%     {'trans_prob_safe_threat_31'       }      0.19702    {0×0 double}       0.29671
%     {'trans_prob_safe_threat_12'       }      0.11019    {0×0 double}       0.56215
%     {'trans_prob_safe_threat_22'       }      0.81517    {'***'     }    4.1301e-08
%     {'trans_prob_safe_threat_32'       }      0.49757    {'**'      }     0.0051462
%     {'trans_prob_safe_threat_13'       }      0.16379    {0×0 double}       0.38712
%     {'trans_prob_safe_threat_23'       }      0.74732    {'***'     }    2.0868e-06
%     {'trans_prob_safe_threat_33'       }      0.66142    {'***'     }    6.9009e-05
%     {'steady_state_prob_safe_threat_1' }      0.68234    {'***'     }    3.2766e-05
%     {'steady_state_prob_safe_threat_2' }      0.75696    {'***'     }    1.2919e-06
%     {'steady_state_prob_safe_threat_3' }       0.7906    {'***'     }    2.0133e-07
%     {'mean_time'                       }      0.62595    {'***'     }    0.00021588
%     {'variance_time'                   }      0.36673    {'*'       }       0.04622
%     {'trans_prob_time_11'              }      0.44695    {'*'       }      0.013279
%     {'trans_prob_time_21'              }      0.80744    {'***'     }    6.9637e-08
%     {'trans_prob_time_31'              }    -0.074725    {0×0 double}       0.69473
%     {'trans_prob_time_12'              }      0.31361    {0×0 double}      0.091488
%     {'trans_prob_time_22'              }      0.80978    {'***'     }    5.9574e-08
%     {'trans_prob_time_32'              }     -0.26272    {0×0 double}       0.16072
%     {'trans_prob_time_13'              }      0.46595    {'**'      }      0.009456
%     {'trans_prob_time_23'              }      0.75472    {'***'     }     1.447e-06
%     {'trans_prob_time_33'              }     0.023468    {0×0 double}       0.90203
%     {'steady_state_prob_time_1'        }      0.81234    {'***'     }    5.0149e-08
%     {'steady_state_prob_time_2'        }      0.78839    {'***'     }    2.2976e-07
%     {'steady_state_prob_time_3'        }       0.4672    {'**'      }     0.0092407
%     {'mean_self-related'               }      0.60899    {'***'     }    0.00035506
%     {'variance_self-related'           }      0.69325    {'***'     }    2.1684e-05
%     {'trans_prob_self-related_11'      }      0.67124    {'***'     }    4.8998e-05
%     {'trans_prob_self-related_21'      }      0.44028    {'*'       }      0.014898
%     {'trans_prob_self-related_12'      }      0.67124    {'***'     }    4.8998e-05
%     {'trans_prob_self-related_22'      }      0.44028    {'*'       }      0.014898
%     {'steady_state_prob_self-related_1'}      0.56528    {'**'      }     0.0011339
%     {'steady_state_prob_self-related_2'}      0.56528    {'**'      }     0.0011339
%     {'mean_vividness'                  }      0.67687    {'***'     }    4.0032e-05
%     {'variance_vividness'              }      0.62129    {'***'     }    0.00024822
%     {'trans_prob_vividness_11'         }      0.59563    {'***'     }    0.00051533
%     {'trans_prob_vividness_21'         }      0.52165    {'**'      }     0.0031126
%     {'trans_prob_vividness_12'         }      0.59563    {'***'     }    0.00051533
%     {'trans_prob_vividness_22'         }      0.52165    {'**'      }     0.0031126
%     {'steady_state_prob_vividness_1'   }      0.59069    {'***'     }    0.00058903
%     {'steady_state_prob_vividness_2'   }      0.59069    {'***'     }    0.00058903

%% TABLE_S4: Correlation between general negative affectivity and Markov-chain dynamic features
for i = 1:size(feature_trained,1)
    [feat_factor_cor(i,1), feat_factor_p(i,1)] = corr(feature_trained(i,:)', F_neg.sess1(:,1));
end
asterisk = pnum2star(feat_factor_p(ind));
t = table(features_descrip_all(ind), feat_factor_cor(ind), asterisk, feat_factor_p(ind), 'VariableNames',{'features','r','asterisk','p'})

%   58×4 table
%                   features                      r           asterisk          p     
%     ____________________________________    __________    ____________    __________
% 
%     {'mean_valence'                    }      -0.43238    {'***'     }    0.00044899
%     {'variance_valence'                }       0.34346    {'**'      }     0.0062755
%     {'trans_prob_valence_11'           }        0.2942    {'*'       }      0.020286
%     {'trans_prob_valence_21'           }       0.35961    {'**'      }     0.0040958
%     {'trans_prob_valence_31'           }       0.41644    {'***'     }    0.00076097
%     {'trans_prob_valence_12'           }      -0.17268    {0×0 double}       0.17955
%     {'trans_prob_valence_22'           }     -0.044684    {0×0 double}       0.73021
%     {'trans_prob_valence_32'           }      -0.20727    {0×0 double}         0.106
%     {'trans_prob_valence_13'           }      -0.22104    {0×0 double}       0.08426
%     {'trans_prob_valence_23'           }       -0.2286    {0×0 double}      0.073917
%     {'trans_prob_valence_33'           }      -0.18034    {0×0 double}       0.16072
%     {'steady_state_prob_valence_1'     }       0.48356    {'***'     }     6.851e-05
%     {'steady_state_prob_valence_2'     }      -0.19307    {0×0 double}       0.13271
%     {'steady_state_prob_valence_3'     }      -0.23155    {0×0 double}      0.070166
%     {'mean_safe_threat'                }      -0.26501    {'*'       }       0.03738
%     {'variance_safe_threat'            }       0.33694    {'**'      }     0.0074087
%     {'trans_prob_safe_threat_11'       }       0.20224    {0×0 double}       0.11494
%     {'trans_prob_safe_threat_21'       }        0.3991    {'**'      }     0.0013125
%     {'trans_prob_safe_threat_31'       }       0.34394    {'**'      }     0.0061978
%     {'trans_prob_safe_threat_12'       }       -0.1192    {0×0 double}       0.35611
%     {'trans_prob_safe_threat_22'       }      -0.20003    {0×0 double}       0.11906
%     {'trans_prob_safe_threat_32'       }      -0.15571    {0×0 double}       0.22685
%     {'trans_prob_safe_threat_13'       }      -0.10861    {0×0 double}       0.40077
%     {'trans_prob_safe_threat_23'       }     -0.029679    {0×0 double}       0.81888
%     {'trans_prob_safe_threat_33'       }     -0.068378    {0×0 double}       0.59745
%     {'steady_state_prob_safe_threat_1' }       0.41495    {'***'     }    0.00079838
%     {'steady_state_prob_safe_threat_2' }      -0.24887    {0×0 double}      0.051108
%     {'steady_state_prob_safe_threat_3' }    -0.0010993    {0×0 double}       0.99323
%     {'mean_time'                       }       0.13709    {0×0 double}       0.28801
%     {'variance_time'                   }       0.15884    {0×0 double}       0.21752
%     {'trans_prob_time_11'              }     0.0051689    {0×0 double}        0.9682
%     {'trans_prob_time_21'              }     0.0037896    {0×0 double}       0.97668
%     {'trans_prob_time_31'              }       -0.1237    {0×0 double}       0.33813
%     {'trans_prob_time_12'              }     -0.009243    {0×0 double}       0.94316
%     {'trans_prob_time_22'              }     -0.056937    {0×0 double}       0.66026
%     {'trans_prob_time_32'              }    -0.0049205    {0×0 double}       0.96972
%     {'trans_prob_time_13'              }      0.011189    {0×0 double}       0.93122
%     {'trans_prob_time_23'              }       0.11362    {0×0 double}       0.37925
%     {'trans_prob_time_33'              }       0.13677    {0×0 double}       0.28915
%     {'steady_state_prob_time_1'        }     -0.027297    {0×0 double}        0.8332
%     {'steady_state_prob_time_2'        }     -0.027953    {0×0 double}       0.82925
%     {'steady_state_prob_time_3'        }       0.10746    {0×0 double}       0.40579
%     {'mean_self-related'               }       0.34802    {'**'      }     0.0055753
%     {'variance_self-related'           }     0.0048486    {0×0 double}       0.97017
%     {'trans_prob_self-related_11'      }       -0.2132    {0×0 double}      0.096153
%     {'trans_prob_self-related_21'      }      -0.25846    {'*'       }       0.04253
%     {'trans_prob_self-related_12'      }        0.2132    {0×0 double}      0.096153
%     {'trans_prob_self-related_22'      }       0.25846    {'*'       }       0.04253
%     {'steady_state_prob_self-related_1'}      -0.27773    {'*'       }      0.028848
%     {'steady_state_prob_self-related_2'}       0.27773    {'*'       }      0.028848
%     {'mean_vividness'                  }          0.27    {'*'       }      0.033813
%     {'variance_vividness'              }       0.14013    {0×0 double}       0.27733
%     {'trans_prob_vividness_11'         }      -0.13262    {0×0 double}       0.30417
%     {'trans_prob_vividness_21'         }      -0.13294    {0×0 double}       0.30297
%     {'trans_prob_vividness_12'         }       0.13262    {0×0 double}       0.30417
%     {'trans_prob_vividness_22'         }       0.13294    {0×0 double}       0.30297
%     {'steady_state_prob_vividness_1'   }      -0.18561    {0×0 double}       0.14864
%     {'steady_state_prob_vividness_2'   }       0.18561    {0×0 double}       0.14864

%% Preparing sets of features and negative affectivity (factor score) for testing


% >>> general negative affectivity >>>
%   fMRI SESS2 data 
q_negativity.sess2 = q_dat(first_idx_n92,:);
[lambda2,psi2,T2,~,F_neg.sess2] = factoran(scale(q_negativity.sess2), 2, 'rotate', 'pattern', 'target', Tgt); % first sess personality
y_actual{1} = F_neg.sess2(:,1);                            
%   WEB SESS1 data                         
y_actual{2} = web.F(:,1);   
%   WEB SESS2 data                  
y_actual{3} = web.F2(:,1);  


% >>> markov chain feature >>>
%   fMRI SESS2 data 
feature_test{1} = ft_sess2';

%   WEB SESS1 data                         
clear x3; k = 0;
for j = trainid_web
    k = k + 1;    feature_temp = [];
    for dim_i = 1:3
        a = web.features{dim_i}{j}';    feature_temp = [feature_temp; a(:)]; 
    end
    x3(:,k) = feature_temp;       
end
feature_test{2} = zeros(58, numel(trainid_web));
feature_test{2}([1:14, 29:42, 43:50],:) = x3([1:14, 23:36, 15:22],:);
feature_test{2} = feature_test{2}';

%   WEB SESS2 data                  
clear x4; k = 0;
for j = testid_web
    k = k + 1;    x_temp = [];
    for dim_i = 1:3
        a = web.features{dim_i}{j}';    x_temp = [x_temp; a(:)]; 
    end
    x4(:,k) = x_temp;       
end
feature_test{3} = zeros(58, numel(testid_web));
feature_test{3}([1:14, 29:42, 43:50],:) = x4([1:14, 23:36, 15:22],:);
feature_test{3} = feature_test{3}';

%% Optimizing models for general negativity

% dat_pred = fmri_data;
% dat_pred.dat = zscore(feature_trained')';
% 
% % predicting factor scores using features
% k = 0;
% dat_pred.Y = F_neg.sess1(:,1);
% for lasso_i = 2:20
%     k = k + 1;
%     [cverr, stats_pred, optout] = predict(dat_pred, 'algorithm_name', 'cv_lassopcr', 'nfolds', 1:numel(dat_pred.Y), 'error_type', 'mse', 'nopcr', 'lasso_num', lasso_i);
%     por(k,1) = stats_pred.pred_outcome_r;
% end
% 
% plot(por)
% 
% lasso_i = 2:20;
% [max_r,max_idx] = max(por);
% 
% lasso_i_max = lasso_i(max_idx);
% lasso_i_max: 12; max_r =  0.5086

%% FINAL MODEL with the best lasso num

lasso_i_max = 12; % to run from here
dat_pred = fmri_data;
dat_pred.dat = zscore(feature_trained')';
dat_pred.Y = F_neg.sess1(:,1);
[cverr, stats, optout] = predict(dat_pred, 'algorithm_name', 'cv_lassopcr', 'nfolds', 1:numel(dat_pred.Y), 'error_type', 'mse', 'nopcr', 'lasso_num', lasso_i_max);

%% SAVE
% datadir = '/Users/byeolkim/dropbox/FAST/sync/data/public sharing';
% savename = fullfile(datadir, 'figure2_b.mat');
% save(savename, 'stats', 'feature_test', 'y_actual', 'features_descrip_all');
