#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun  5 16:02:23 2019

This is a demo code of how we used sklearn toolbox to perform the 
multivariate analyses . 

@author: podvae01
"""
import HLTP # data/results directories definitions
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import LeaveOneOut, cross_val_predict
from sklearn.metrics import roc_auc_score
from mne.decoding import LinearModel, get_coef

def Prepare_Data_and_Labels(decoded_variable, trial_subset, label_values):
    '''A function loading the data and labels, 
    i.e., data recorded in a subset of trials with their labels 
    according to the variable that we want to decode'''
    # Load all neural and behaviorl data
    # Select subset of trials (trial_subset)
    # Reject trials that do not correspond to label_values (eg. missing report)
    return data, labels

# --------------------- model definition (GENERAL) ---------------------------
model = make_pipeline(StandardScaler(), 
                      LogisticRegression(C = 1, solver='liblinear'))
cv = LeaveOneOut()  #Cross-validation method

# load data 
decoded_variable = 'recognition'
data, labels = Prepare_Data_and_Labels(
        decoded_variable, 
        trial_subset = {'real_img':True}, 
        label_values = [1, -1])
# MVPA for each subject 
pred_prob = []
for sub_idx in enumerate(HLTP.subjects): 
    # using  cross-validation to calculate probabiltiy of the two labels 
    # for each left-out-trial
    pred_prob.append(cross_val_predict(model, data[sub_idx], 
                labels[sub_idx], cv = cv, 
                method = 'predict_proba', n_jobs = -1))
    
# --------------------- model definition (SPECIFIC) --------------------------
model = make_pipeline(StandardScaler(), LogisticRegression(
                                      C = 1, multi_class = 'multinomial',
                                      solver = 'newton-cg'))
cv = LeaveOneOut()

decoded_variable = 'category'
data, labels = Prepare_Data_and_Labels(
        decoded_variable, {'real_img':True}, label_values = [1, 2, 3, 4])

pred_prob = []
for sub_idx in enumerate(HLTP.subjects): 
    # using  cross-validation to calculate probabiltiy of the four labels 
    # for each left-out-trial neural data trial
    pred_prob.append(cross_val_predict(model, data[sub_idx], 
                                labels[sub_idx], cv = cv, 
                                method='predict_proba', n_jobs = -1))

# ------ Example of how to fit the model in one set and test in abother -------
    
data_train, labels_train = Prepare_Data_and_Labels(decoded_variable, 
        trial_subset = {'real_img':True}, label_values = [1, -1])    
data_test, labels_test = Prepare_Data_and_Labels(decoded_variable, 
        trial_subset = {'real_img':False}, label_values = [1, -1]) 

model.fit(data_train[sub_idx], labels_train[sub_idx])
pred_prob = model.predict_proba(data_test[sub_idx])

# calculate AUROC:
auroc = roc_auc_score(labels_test, pred_prob) 

# calculate activation patterns:
model = make_pipeline(StandardScaler(), LinearModel(LogisticRegression(
                                      C = 1, multi_class = 'multinomial',
                                      solver = 'newton-cg'))) 
model.fit(data_train[sub_idx], labels_train[sub_idx])
patterns = get_coef(model, 'patterns_', inverse_transform = True)