from mvpa.suite import *
from mvpa.mappers.procrustean import ProcrusteanMapper
from mvpa.algorithms.hyperalignment import Hyperalignment

mask_suffix = '_mask_vt_new.nii.gz'
movie_data_suffix = '_movie.nii.gz'
mkdg_data_suffix = '_mkdg.nii.gz'
''' Number of features/voxels to use from each hemisphere'''
nf = 500

''' Subject IDs '''
subj = ['cb', 'dm', 'hj', 'kd', 'kl', 'mh', 'ph', 'rb', 'se', 'sm']
mappers_all = []
print "Loading movie data..."
'''
Loading movie data of subjects and adding mask information as 
feature attribute 

'''
ds = []
for sub in subj: 
    ds.append(fmri_dataset(samples=sub+movie_data_suffix,mask=sub+mask_suffix, add_fa={'hem':sub+mask_suffix}))

ds = [vstack([ds_[4:1101,:], ds_[1101+5:2212,:]]) for ds_ in ds]

ds_full = deepcopy(ds)

'''
Compute feature ranks in each dataset
based on correlation with other subjects' datasets

'''
print "Computing feature ranks..."
fs = [];
for hem in xrange(2):
    if hem == 0:
        ds = [sd[:, sd.fa.hem == 0.99993917097890517 ] for sd in ds_full]
    else:
        ds = [sd[:, sd.fa.hem == 1.9999395040576928 ] for sd in ds_full]
    feature_scores = [ np.zeros(ds[i].nfeatures) for i in xrange(len(subj)) ]
    for i, sd in enumerate(ds):
        ds_temp = sd.copy()
        zscore(ds_temp, chunks_attr=None)
        for j, sd2 in enumerate(ds[i+1:]):
            ds_temp2 = sd2.copy()
            zscore(ds_temp2, chunks_attr=None)
            corr_temp= np.asarray(np.mat(np.transpose(ds_temp.samples))*np.mat(ds_temp2.samples))
            feature_scores[i] = feature_scores[i] + np.max(corr_temp, axis = 1)
            feature_scores[j+i+1] = feature_scores[j+i+1] + np.max(corr_temp, axis = 0)
  
    fs.append(feature_scores)

'''
Computing hyperalignment between subjects'
left and right hemispheres separately

'''
print "Computing transformations for "
print nf
fselector = FixedNElementTailSelector(nf, tail='upper', mode='select',sort=False)
mapper_results = [[]]*2
for hem in xrange(2):
    if hem == 0:
        ds = [sd[:,sd.fa.hem == 0.99993917097890517 ] for sd in ds_full]
    else:
        ds = [sd[:,sd.fa.hem == 1.9999395040576928 ] for sd in ds_full]
  
    for i in xrange(len(ds)):
        ds[i].fa['bsc_scores'] = fs[hem][i]
  
    ds_fs = [ sd[:, fselector( sd.fa.bsc_scores )] for sd in ds]
    hyper = Hyperalignment(zscore_common=True, ref_ds = 0)
    mapper_results[hem] = hyper(datasets=ds_fs)
  
mappers_all.append(mapper_results)
print "Hyperalignment parameters computed"

'''

Validation of hyperalignment on monkeydog dataset
by performing between subject classification

'''

md_cd = ColumnData('labels.txt',header=['label'])
md_labels = [int(x) for x in md_cd['label']]

# Making run ends as -1s 
for run in range(8):
    md_labels[192*run:192*run+3] = [-1]*3
print "Loading monkeydog data..."
# Loading monkeydog data of subjects
mkdg_ds = []
for sub in subj:
    mkdg_ds.append(fmri_dataset(samples=sub+mkdg_suffix,targets=md_labels, chunks=np.repeat(range(8),192) , mask=sub+mask_suffix,add_fa={'hem':sub+mask_suffix}))

m=mean_group_sample(['targets', 'chunks'])

# Averaging blocks of categories
mkdg_ds = [mkdg_ds[i].get_mapped(m) for i in range(len(mkdg_ds))]
mkdg_ds = [mkdg_ds[i][mkdg_ds[i].sa.targets != -1] for i in range(len(mkdg_ds))]
mkdg_ds_full = deepcopy(mkdg_ds)

bsc_accuracies = []
wsc_accuracies = []
wsc_hyper_accuracies = []

clf = clfswh['multiclass','linear','NU_SVC'][0]
cvterr = CrossValidation(clf, NFoldPartitioner())
terr = TransferMeasure(clf, Splitter(attr='bsc_split',attr_values=[0, 1]), postproc= lambda x:float(np.sum(np.equal(x.samples.T,x.targets)))/x.nsamples )
print "Performing classifications using"
print nf
fselector = FixedNElementTailSelector(nf, tail='upper', mode='select',sort=False)
mkdg_data = []
for hem in xrange(2):
    if hem == 0:
        mkdg_ds =  [sd[:,sd.fa.hem == 0.99993917097890517 ] for sd in mkdg_ds_full]
    else:
        mkdg_ds =  [sd[:,sd.fa.hem == 1.9999395040576928 ] for sd in mkdg_ds_full]
   
    ''' Adding feature scores to datasets '''
    for i in range(len(mkdg_ds)):
        mkdg_ds[i].fa['bsc_scores'] = fs[hem][i]
    if mkdg_data == []:
        mkdg_data = mkdg_ds
    else:
        mkdg_data += mkdg_ds

print "Selecting features..."
mkdg_orig_fs = [ sd[:, fselector( sd.fa.bsc_scores )]  for sd in mkdg_data]

print "Hyperaligning data..."
for ds_ in mkdg_orig_fs: zscore(ds_,chunks_attr=None) 

mkdg_hyper  = [ mappers_all[inf][int(i>9)][ i - 10*(i>9)].forward(ds_) for i,ds_ in enumerate(mkdg_orig_fs)]

mkdg_ds_fs = []
mkdg_ds_mapped = []
for i in range(10):
    mkdg_ds_fs.append( Dataset( samples= np.hstack((mkdg_orig_fs[i], mkdg_orig_fs[i+10])), sa={'targets': mkdg_orig_fs[i].targets, 'chunks': mkdg_orig_fs[i].chunks} ) )
    mkdg_ds_mapped.append( Dataset( samples= np.hstack((mkdg_hyper[i], mkdg_hyper[i+10])), sa={'targets': mkdg_hyper[i].targets, 'chunks': mkdg_hyper[i].chunks} ) )

print "Zscoring and removing blanks"
for ds_ in mkdg_ds_fs: zscore(ds_, chunks_attr=None)
mkdg_ds_fs = [sd[sd.sa.targets != 0] for sd in mkdg_ds_fs]

# within-subject classification
print "Performing wsc..."
for sd in mkdg_ds_fs:
    wsc = cvterr(sd)
    wsc_accuracies.append(1-np.mean(wsc))

print "Performing wsc hyper..."
for ds_ in mkdg_ds_mapped: zscore(ds_, chunks_attr=None)

mkdg_ds_mapped = [sd[sd.sa.targets != 0] for sd in mkdg_ds_mapped]

for sd in mkdg_ds_mapped:
    wsc = cvterr(sd)
    wsc_hyper_accuracies.append(1-np.mean(wsc))

print "Performing bsc..."
mkdg_ds_all = vstack(mkdg_ds_mapped)
mkdg_ds_all.sa['subject'] = np.repeat(range(10),56)
mkdg_ds_all.sa['runs'] = mkdg_ds_all.sa['chunks']
mkdg_ds_all.sa['chunks'] = mkdg_ds_all.sa['subject']
mkdg_ds_all.sa['bsc_split'] = np.zeros(80*7,dtype=np.int)
for test_run in range(8):
    for test_subj in range(10):
        mkdg_ds_train = mkdg_ds_all[ np.where( np.logical_and( mkdg_ds_all.sa.runs != test_run, mkdg_ds_all.sa.subject != test_subj)) ]
        mkdg_ds_test  = mkdg_ds_all[ np.where( np.logical_and( mkdg_ds_all.sa.runs == test_run, mkdg_ds_all.sa.subject == test_subj)) ]
        mkdg_ds_train.sa.bsc_split = np.zeros(63*7,dtype=np.int)
        mkdg_ds_test.sa.bsc_split = np.ones(7,dtype=np.int)
        bsc_accuracies.append( terr(vstack([mkdg_ds_train, mkdg_ds_test])))

bsc_accuracies = np.reshape(bsc_accuracies, (8,10))
print "BSC hyperalignment:"
print np.mean(bsc_accuracies)
print "WSC"
print np.mean(wsc_accuracies)