import numpy as np
np.random.seed(3)
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

fmri.GenHRF(PATHS).gen_test_data()

%matplotlib inline
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'}, rois='rh_LO')
mvpa.plot_roi(roi='rh_LO.nii')

%matplotlib inline
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'},
                     rp={'values':'raw', 'method':None, 'visualize':True},
                     rois='rh_LO', dur=4, offset=0)
df, df_fname, loaded = mvpa.run()

%matplotlib inline
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'},
                     rp={'values':'raw', 'method':'timecourse', 'force':'all'},
                     rois='rh_LO', dur=8, offset=0)
df, df_fname, loaded = mvpa.run()

%matplotlib inline
from psychopy_ext import exp, fmri
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style='dark')
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'},
                     rp={'values':'raw', 'method':None, 'force':'all'},
                     rois='rh_LO', dur=2, offset=2)
ds = mvpa.extract_samples('subj_01', 'main', ['rh_LO.nii', 'rh_LO', '*rh_LO.nii'],
                          values='raw', offset=2)
ds = mvpa.detrend(ds)
evds = mvpa.ds2evds(ds, dur=2)
mvpa.plot_ds(ds)

import mvpa2.suite
import numpy as np

# leave fixation out
evds = evds[evds.sa.targets != mvpa.fix]

# calculate the mean per target per chunk (across trials)
run_averager = mvpa2.suite.mean_group_sample(['targets','chunks'])
evds_avg = evds.get_mapped(run_averager)
numt = len(evds_avg.UT)

# calculate mean across conditions per chunk per voxel
target_averager = mvpa2.suite.mean_group_sample(['chunks'])
mean = evds_avg.get_mapped(target_averager)
# subtract the mean chunk-wise
evds_avg.samples -= np.repeat(mean, numt, 0)

targets = evds_avg.UT

# split 1
evds_split1 = evds_avg[:4]
run_averager = mvpa2.suite.mean_group_sample(['targets'])
evds_split1 = evds_split1.get_mapped(run_averager)

# split 2
evds_split2 = evds_avg[4:]
run_averager = mvpa2.suite.mean_group_sample(['targets'])
evds_split2 = evds_split2.get_mapped(run_averager)

result = mvpa2.clfs.distance.one_minus_correlation(evds_split1.samples,
    evds_split2.samples) / 2

plt.imshow(result, cmap='coolwarm', interpolation='none')
plt.colorbar()
plt.show()

# plot averages for within-condition and across-conditions
within = result.trace() / result.shape[0]
across = (np.sum(result) - result.trace()) / (result.size - result.shape[0])
ax = plt.subplot(111)
ax.bar(0, within, label='within')
ax.bar(1, across, label='across')
ax.legend()
plt.show()

%matplotlib inline
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'},
                     rp={'values':'raw', 'method':'corr', 'plot':True},
                     rois='rh_LO', dur=2, offset=2)
df, df_fname, loaded = mvpa.run()

"""
Modified from PyMVPA example http://www.pymvpa.org/examples/pylab_2d.html
and http://www.pymvpa.org/examples/pylab_2d.html
"""
import mvpa2.suite
import numpy as np
import matplotlib.pyplot as plt
import seaborn
%matplotlib inline

np.random.seed(3)

##############
# Parameters #
##############
npoints = 200
xymax = 9  # Absolute max value allowed. Just to assure proper plots
ngrid = 101  # grid for evaluting performance

# set grid
x1 = np.linspace(0, xymax, ngrid)
x2 = np.linspace(0, xymax, ngrid)
x,y = np.meshgrid(x1, x2)
feat_grid = np.array((np.ravel(x), np.ravel(y)))

def train_svm(ds, cl= mvpa2.suite.LinearNuSVMC):
    # define classifier
    clf = cl(probability=1, enable_ca=['probabilities'])
    # enable saving of the estimates used for the prediction
    clf.ca.enable('estimates')
    # train with the known points
    clf.train(ds)
    # run the predictions on the test values
    pre = clf.predict(feat_grid.T)
    # get probabilities from the svm
    res = np.asarray([(q[1][1] - q[1][0] + 1) / 2
                for q in clf.ca.probabilities])
    # get support vectors
    supvecs = clf.model.get_sv()
    
    return res, supvecs

def plot_data(ds, supvecs=None, contour=None):
    ax = plt.subplot(111)
    ax.set_aspect('equal')
    
    # plot samples
    ax.plot(ds.samples[ds.targets == 1, 0], ds.samples[ds.targets == 1, 1], '.',
            color='indianred', label='condition 1')
    ax.plot(ds.samples[ds.targets == 0, 0], ds.samples[ds.targets == 0, 1], '.',
            color='steelblue', label='condition 2')    
    
    # plot support vectors as larger dots
    if supvecs is not None:
        feat_pos = ds.samples[ds.targets == 1]
        feat_neg = ds.samples[ds.targets == 0]
        pos_test = np.array([f for f in feat_pos if f in supvecs]).T
        neg_test = np.array([f for f in feat_neg if f in supvecs]).T
        ax.plot(pos_test[0], pos_test[1],
                'o', color='indianred')
        ax.plot(neg_test[0], neg_test[1],
                'o', color='steelblue')
    
    # plot probability contour
    if contour is not None:
        # plot decision surfaces at few levels to emphasize the topology
        z = np.asarray(contour).reshape((ngrid, ngrid))
        
        ax.contour(x, y, z, [0.1, 0.4, 0.5, 0.6, 0.9],
                   linestyles=['dotted', 'dashed', 'solid', 'dashed', 'dotted'],
                   linewidths=1, colors='black', hold=True)
        
    ax.set_xlabel('response in voxel 1')
    ax.set_ylabel('response in voxel 2')
    ax.legend()

#################
# Generate data #
#################
dist = 1  # distance between the two distributions

def gen_data(dist):
    feat = mvpa2.suite.pure_multivariate_signal(30, 2)
    feat = np.random.randn(2, npoints) + dist + xymax / 2.
    feat = feat.clip(0, xymax)
    return feat

# two categories of samples
feat_pos = gen_data(dist)
feat_neg = gen_data(-dist)

# create the pymvpa dataset from the labeled features
pat_pos = mvpa2.suite.dataset_wizard(samples=feat_pos.T, targets=1)
pat_neg = mvpa2.suite.dataset_wizard(samples=feat_neg.T, targets=0)
ds = mvpa2.suite.vstack((pat_pos, pat_neg))

# plot data
plot_data(ds)

# train SVM and plot
res, supvecs = train_svm(ds)
plot_data(ds, supvecs=supvecs, contour=res)

%matplotlib inline
import numpy as np
np.random.seed(3)
from psychopy_ext import exp, fmri
PATHS = exp.set_paths(fmri_rel='%s')

mvpa = fmri.Analysis(PATHS, 2, info={'subjid':'subj_01', 'runtype':'main'},
                     rp={'values':'raw', 'method':'svm', 'plot':True},
                     rois='rh_LO', dur=2, offset=2)
df, df_fname, loaded = mvpa.run()

ds = mvpa2.suite.pure_multivariate_signal(npoints, 2)
ds.samples += 6
plot_data(ds)

res, supvecs = train_svm(ds, cl=mvpa2.suite.RbfNuSVMC)
plot_data(ds, supvecs=supvecs, contour=res)