""" Tests the different linear regression models 
and chooses best model using AIC, BIC, and Adjusted R2

Averages the criteria for all voxels across 3 subjects

Run with: 
    python model_selection.py
"""

from __future__ import absolute_import, division, print_function
import numpy as np
import numpy.linalg as npl
import nibabel as nib
import pandas as pd # new
import sys # instead of os
import os


# Relative path to subjects
project_path          = "../../../"
path_to_data          = project_path+"data/ds009/"
location_of_images    = project_path+"images/"
location_of_functions = project_path+"code/utils/functions/" 
final_data            = "../../../final/data/"
smooth_data           =  final_data + 'smooth/'
hrf_data              = final_data + 'hrf/'
behav_suffix           = "/behav/task001_run001/behavdata.txt"


sys.path.append(location_of_functions)
sub_list = os.listdir(path_to_data)
sub_list = [i for i in sub_list if 'sub' in i]


from event_related_fMRI_functions import hrf_single, np_convolve_30_cuts
from time_shift import time_shift, make_shift_matrix, time_correct
from glm import glm_multiple, glm_diagnostics
from noise_correction import mean_underlying_noise, fourier_predict_underlying_noise,fourier_creation
from Image_Visualizing import present_3d, make_mask
from mask_phase_2_dimension_change import masking_reshape_start, masking_reshape_end
from model_comparison import adjR2, BIC, AIC
from hypothesis import t_stat_mult_regression



#Choose which criteria
input_var = 'AIC'

##################################
# Functions for different models #
##################################

if input_var == 'adjR2':
    
    def model(MRSS,y_1d,df, rank):
	    return adjR2(MRSS,y_1d, df, rank)
        
elif input_var == 'BIC': 
       
    def model(MRSS,y_1d,df, rank):
        return BIC(MRSS, y_1d, df, rank)
        
elif input_var == 'AIC': 
       
    def model(MRSS,y_1d,df, rank):
        return AIC(MRSS, y_1d, df, rank)
    

#List to include all of the models
model1=[]
model2=[]
model3=[]
model4=[]
model4_5=[]
model5=[]
model6=[]
model7=[]
model8=[]
model9=[]
model9_5=[]
model10=[]


#Average for subject 2,3,and 14
for i in ['sub002','sub003','sub014']:
    img = nib.load(smooth_data+ i +"_bold_smoothed.nii")
    data = img.get_data()

    behav=pd.read_table(path_to_data+i+behav_suffix,sep=" ")
    num_TR = float(behav["NumTRs"])


    #CREATE THE CONVOLVE STUFF
    cond1=np.loadtxt(path_to_data+ i+ "/model/model001/onsets/task001_run001/cond001.txt")
    cond2=np.loadtxt(path_to_data+ i+ "/model/model001/onsets/task001_run001/cond002.txt")
    cond3=np.loadtxt(path_to_data+ i+ "/model/model001/onsets/task001_run001/cond003.txt")

    TR = 2
    tr_times = np.arange(0, 30, TR)
    hrf_at_trs = np.array([hrf_single(x) for x in tr_times])

    # creating the .txt file for the events2neural function
    cond_all=np.row_stack((cond1,cond2,cond3))
    cond_all=sorted(cond_all,key= lambda x:x[0])

    cond_all=np.array(cond_all)[:,0]

    delta_y=2*(np.arange(34))/34

    shifted_all=make_shift_matrix(cond_all,delta_y)
    shifted_1= make_shift_matrix(cond1[:,0],delta_y)
    shifted_2= make_shift_matrix(cond2[:,0],delta_y)
    shifted_3= make_shift_matrix(cond3[:,0],delta_y)

    #########################################
    #Create convovled HRF for each condition#
    #########################################

    def make_convolve_lambda(hrf_function,TR,num_TRs):
        convolve_lambda=lambda x: np_convolve_30_cuts(x,np.ones(x.shape[0]),hrf_function,TR,np.linspace(0,(num_TRs-1)*TR,num_TRs),15)

        return convolve_lambda

    convolve_lambda=make_convolve_lambda(hrf_single,TR,num_TR)

    hrf_matrix_all=time_correct(convolve_lambda,shifted_all,num_TR)
    hrf_matrix_1=time_correct(convolve_lambda,shifted_1,num_TR)
    hrf_matrix_2=time_correct(convolve_lambda,shifted_2,num_TR)
    hrf_matrix_3=time_correct(convolve_lambda,shifted_3,num_TR)

    n_vols = data.shape[-1]

    ##########################################
    # Create PCA features for our regression #
    ##########################################
    mask = nib.load(path_to_data+i+'/anatomy/inplane001_brain_mask.nii.gz')
    mask_data = mask.get_data()