import scipy as sp
from scipy import stats

import nipy.modalities.fmri.design_matrix as dm
from nipy.modalities.fmri.experimental_paradigm import EventRelatedParadigm

# import rpy
from rpy2.robjects import r
from rpy2.robjects import FloatVector, IntVector, ListVector

# Import neuRosim
r['library']('neuRosim')

def simulate_data(onsets, # Onsets per condition, per region
                  nscans=100, # Number of scans
                  TR=2.0,     # Tr
                  region_centers = [[20, 20], [60, 60]],  # Centers of regions
                  region_radiuses = [15, 10],             # Size of regions
                  effect_sizes = [[2.0], [.75]],        # Effect sizes of conditions per region per 
                                                          # condition [[r1c1, r1c2], [r2c1, r2c2]]
                  durations = [[2.0], [2.0]],             # Durations per region, per condition
                  SNR = 15,       # SNR
                  FWHM = 5,        # Amount of smoothing                        
                  n_regions = 2,   # Number of regions
                  fading = 0.1,    # Amount of fading from regions to outside
                  size = [80, 80], # Total size of images
                  ):
       
    # Get total duration of scan
    total_time = nscans * TR
    
    # Convert input to R objects
    onsets = [[FloatVector(ol_cond) for ol_cond in ol_region] for ol_region in onsets]
    region_centers = [FloatVector(rc) for rc in region_centers]
    size = IntVector(size)
    

    # Set up design (neuRosim)
    design = r['simprepTemporal'](regions=n_regions, 
                                  onsets=onsets, 
                                  durations=durations, 
                                  TR=TR, 
                                  totaltime=total_time, 
                                  effectsize=effect_sizes)

    
    # Set up the spatial map (neuRosim)
    spatial = r['simprepSpatial'](regions=n_regions, 
                                  coord=region_centers, 
                                  radius=region_radiuses,
                                  form="sphere", 
                                  fading=fading)
    
    # Get the simulated tata (neuRosim)
    sim_data = r['simVOLfmri'](design=design, 
                               image=spatial, 
                               SNR=SNR, 
                               noise="mixture", 
                               dim=size, 
                               w=FloatVector([0.05,0.1,0.01,0.09,0.05,0.7]),
                               FWHM=FWHM,
                               spat='gammaRF')
                    
               

    
    return np.array(sim_data)



def simulate_subject(spatial_jitter, 
                     radius_jitter, 
                     effect_size_jitter, 
                     effect_sizes=[2.5, 0.75],
                     radiuses=[15, 10],
                     snr_jitter=1, 
                     region_centers=[[25, 25], [55, 55]],
                     standard_snr=2.0):
    
    # Get region centers 
    region_centers =sp.stats.norm(loc=region_centers, scale=spatial_jitter).rvs((2,2))
    
    
    # Get radiuses
    radiuses = list(sp.stats.norm(loc=radiuses, scale=radius_jitter).rvs(2))
    
    # Get effect sizes
    effect_sizes = sp.stats.truncnorm(0, np.inf, loc=effect_sizes, scale=effect_size_jitter).rvs(2)
    effect_sizes = [[e] for e in list(effect_sizes)]
    
    # Get SNR
    snr = sp.stats.norm(loc=standard_snr, scale=snr_jitter).rvs(1)
    
    return region_centers, radiuses, effect_sizes, snr

n_subjects = 20
n_scans = 500
TR = 2.0
onsets = np.arange(10, n_scans*TR, 15)                                                  


region_centers, radiuses, effect_sizes, snr = zip(*[simulate_subject(5, 5, .5) for i in np.arange(n_subjects)])

ms_data = [] #Multiple subjects-data...

for (region_center, radius, effect_size, snr) in zip(region_centers, radiuses, effect_sizes, snr):
    ms_data.append(simulate_data([[onsets]] *2, # For both regions, for one condition...
                 nscans=n_scans,
                 FWHM=5,
                 fading=0.01,
                 region_centers=region_center,
                 region_radiuses=radius,
                 effect_sizes=effect_size,))

np.save('/home/gdholla1/notebooks/ms_data.np', ms_data)

ms_data = np.load('/home/gdholla1/notebooks/ms_data.np.npy')

def fit_sandwich_glm(data, onsets, q=25, c=None):
    
    if c == None:
        c = np.array([[1, 0, 0]])
    
    # Set up paradigm and design matrix
    paradigm = EventRelatedParadigm(['task']*len(onsets), onsets)
    TR = 2.0
    frametimes = np.arange(0, data.shape[-1] * TR, TR) 

    X, names = dm.dmtx_light(frametimes, paradigm, drift_model='polynomial',
                 hfcut=128, hrf_model='canonical')
    
    
    # cut up data
    cuts = np.linspace(0, data.shape[-1], q+1)

    data_cut = np.array([data[:, :, cuts[i]:cuts[i+1]] for i in np.arange(q)])
    X_cut = np.array([X[cuts[i]:cuts[i+1], :] for i in np.arange(q)])
    
    beta = np.zeros((X.shape[1],)+ data.shape[:-1])  # (num_regressors, size_x, size_y)    
    
    # Fit GLM
    predicted = np.zeros_like(data_cut)
    residuals = np.zeros_like(data_cut)

    for i, x in enumerate(X_cut):
        # tensortdot is necessary to indicate which axes need to be 'summed' and get beta
        beta += np.einsum('ij,klj', np.linalg.pinv(x.T.dot(x)).dot(x.T), data_cut[i])

    beta /= q

    for i, x in enumerate(X_cut):
        # Now we can get predicted timecourses
        predicted[i] = np.tensordot(beta, x, (0, 1))

        # ... and residuals
        residuals[i] = data_cut[i] - predicted[i]
    
    # Estimate covariance residuals
    W = np.zeros((residuals.shape[-1],) * 2 + data.shape[:-1]) # (length_block, length_block, width, height)

    for residual in residuals:
        W += np.einsum('...i, ...j->ij...', residual, residual) # do outer product over last two dimensions, broadcast the rest

    W /= (q - 1)
    
    # Give variance matrix
    V = np.zeros((X.shape[1],) * 2 + data.shape[:-1]) # (num_regresor, num_regressor, width, height)

    for i, x in enumerate(X_cut):

        sandwich1 = np.linalg.pinv(x.T.dot(x)).dot(x.T)    
        sandwich2 = x.dot(np.linalg.pinv(x.T.dot(x)))

        # Apply first part sandwich
        v = np.tensordot(sandwich1, W, axes=((1), (0)))

        # Apply second part
        v = np.tensordot(v, sandwich2, (1, 0))

        # Roll axis to get correct shape
        v = np.rollaxis(v, -1, 1)

        V += v
    
    # Correct for number of replications
    V /= q
    
    # Get t- and z-maps
    dof = q - 1
    t = np.tensordot(c, beta, (1, 0)).squeeze()
    t /= np.tensordot(np.tensordot(c, V, (1,0)), c.T, (1, 0)).squeeze()
    
    z = sp.stats.norm.ppf(sp.stats.t(dof).cdf(t))

    # Correct for numerical rounding errors
    z[z == np.inf] = z[z != np.inf].max()
    
    return beta, V, residuals, t, z, X_cut
    
    

betas, Vs, residualss, ts, zs, design_matrices = zip(*[fit_sandwich_glm(d, onsets) for d in ms_data])

plt.figure(figsize=(15, 10))
for i, z in enumerate(zs):
    plt.subplot(5,4, i +1)
    plt.imshow(z, cmap=plt.cm.hot)
    plt.title('Subject %d' % (i + 1))
    plt.colorbar()
    
plt.tight_layout()

import nibabel as nb

# Degrees of freedom
dof = len(ms_data) - 1
np.savetxt('dof.txt', [dof], fmt='%d')


# Just take a contrast with only first (task-)parameter for now
for i in np.arange(len(ms_data)):
    nb.save(nb.Nifti1Image(betas[i][0][..., np.newaxis], np.identity(4)), 'cope_%d.nii.gz' % i)
    nb.save(nb.Nifti1Image(Vs[i][0, 0][..., np.newaxis], np.identity(4)), 'varcope_%d.nii.gz' % i)
    
# Make a simple mask
nb.save(nb.Nifti1Image(np.ones_like(betas[0][0]), np.identity(4)), 'mask.nii.gz')

import nipype.interfaces.io as nio
import nipype.interfaces.fsl as fsl
import nipype.pipeline.engine as pe
from nipype.workflows.fmri.fsl.estimate import create_fixed_effects_flow
import os

# Standard fixed effets model
level2 = create_fixed_effects_flow()
level2.inputs.flameo.run_mode = 'flame1'
level2.base_dir = '/home/gdholla1/workflow_folders/'


# Use the contrasts and variances from the GLM
level2.inputs.inputspec.copes = [os.path.abspath('cope_%d.nii.gz') % i for i in np.arange(len(ms_data))]
level2.inputs.inputspec.varcopes = [os.path.abspath('varcope_%d.nii.gz') % i for i in np.arange(len(ms_data))]
level2.inputs.inputspec.dof_files = [os.path.abspath('dof.txt')] * len(ms_data)
level2.inputs.l2model.num_copes = len(ms_data)
level2.inputs.flameo.mask_file = os.path.abspath('mask.nii.gz')


# Estimate smoothness for GRF
smooth_est = pe.Node(fsl.SmoothEstimate(), name='smooth_estimate')
smooth_est.inputs.mask_file = os.path.abspath('mask.nii.gz')
pickfirst = lambda x: x[0]
level2.connect(level2.get_node('flameo'), ('zstats', pickfirst), smooth_est, 'zstat_file')

cluster = pe.Node(fsl.Cluster(), name='cluster')
cluster.inputs.threshold = 2.3
cluster.inputs.out_threshold_file  = 'z_simulation_threshold.nii.gz'
level2.connect(level2.get_node('flameo'), ('zstats', pickfirst), cluster, 'in_file')
level2.connect(smooth_est, 'dlh', cluster, 'dlh')

# Write data to datasink
ds = pe.Node(nio.DataSink(), name='datasink')
ds.inputs.base_directory = '/home/gdholla1/data/in_limbo'

level2.connect(level2.get_node('flameo'), ('zstats', pickfirst), ds, 'zstats_simulation')
level2.connect(cluster, 'threshold_file', ds, 'zstats_simulation.@out_threshold_file')

level2.run()

fig = plt.figure(figsize=(14,6))
plt.subplot(121)
z = nb.load('/home/gdholla1/data/in_limbo/zstats_simulation/_flameo0/zstat1.nii.gz').get_data()
plt.imshow(z, cmap=plt.cm.hot, vmin=z.min(), vmax=z.max())
plt.grid(False)
plt.title('raw Level 2 z-map')

plt.subplot(122)
z_threshold = nb.load('/home/gdholla1/data/in_limbo/zstats_simulation/z_simulation_threshold.nii.gz').get_data()
im = plt.imshow(z_threshold, cmap=plt.cm.hot, vmin=z.min(), vmax=z.max())
plt.grid(False)

# Get comparison voxel
comparison_index = np.unravel_index(np.ma.argmin(np.ma.masked_array(z_threshold, z_threshold == 0)), z.shape)
plt.scatter(*comparison_index[::-1], marker='o', s=500, c='r', lw=5, facecolors='none', label='comparison voxel')
plt.xlim(0, 80)
plt.ylim(80, 0)


plt.title('Thresholded level 2 z-map (z>2.3, cluster-based correction)')
# plt.colorbar()

fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
fig.colorbar(im, cax=cbar_ax)



def get_covariance_with_comparison_voxel(residuals, comparison_index, design_matrices):
        
    # Get number of replications
    q = residuals.shape[0]
    
    # Get the residuals in the comparison voxel
    residuals_comparison_voxel = residuals[(Ellipsis,) + comparison_index + (Ellipsis,)]

    # Allocate memory for W_ij matrix
    W_ij = np.zeros((residuals.shape[-1],) * 2 + residuals.shape[1:-1]) # (length_block, length_block, width, height)

    
    # Create covariance matrices for every residual
    for comp_residual, residual in zip(residuals_comparison_voxel, residuals):
        W_ij += np.einsum('...i, ...j->ij...', comp_residual, residual) # do outer product over last two dimensions, broadcast the rest
    
    # Normalize
    W_ij /= (q - 1)
    
    # Allocate memory for the covariance matrix of every voxel with the comparison voxel
    V_ij = np.zeros((design_matrices.shape[-1],) * 2 + residuals.shape[1:-1]) # (num_regresor, num_regressor, width, height)

    for i, x in enumerate(design_matrices):

        sandwich1 = np.linalg.pinv(x.T.dot(x)).dot(x.T)    
        sandwich2 = x.dot(np.linalg.pinv(x.T.dot(x)))

        # Apply first part sandwich
        v = np.tensordot(sandwich1, W_ij, axes=((1), (0)))

        # Apply second part
        v = np.tensordot(v, sandwich2, (1, 0))

        # Roll axis to get correct shape
        v = np.rollaxis(v, -1, 1)
        
        V_ij += v

    V_ij /= q
    
    return V_ij

V_ijs = [get_covariance_with_comparison_voxel(residuals, comparison_index, design_matrix) for residuals, design_matrix in zip(residualss, design_matrices)]
V_ijs = np.array(V_ijs)

Vs = np.array(Vs)

c = np.array([[1, 0, 0]])
V1 = np.tensordot(np.tensordot(c, Vs, (1,1)).squeeze(), c.T, (1, 0))
V2 = np.tensordot(np.tensordot(c, Vs[(Ellipsis, ) + comparison_index], (1,1)), c.T, (-1, 0))
V3 = np.tensordot(np.tensordot(c, V_ijs, (1,1)).squeeze(), c.T, (1, 0))


W_ill = np.zeros((Vs.shape[0], Vs.shape[0]) + ms_data.shape[1:-1])  # Allocate diagonal matrix for every voxel

W_ill[np.arange(Vs.shape[0]), np.arange(Vs.shape[0]), ...] = 1 / np.sqrt((V1.squeeze().T + V2.squeeze() - 2*V3.squeeze().T).T) # Squeeze and transform to facilitate broadcasting

Z = (np.array(betas)[(Ellipsis, 0) + comparison_index] - np.array(betas)[:, 0, ...].T).T # Transform to facilitate broadcasting

W = W_ill[:, :, 20, 20]
Y = Z[:, 20, 20]
G = np.ones((W_ill.shape[0], 1))

import statsmodels.api as sm

beta = pinv(G.T.dot(W).dot(G)).dot(G.T).dot(W).dot(Y)
Y_ = G.dot(beta)
s2 = (Y - Y_).T.dot(W).dot((Y - Y_))/(Y.shape[0] - 1)
wls_V = s2 *pinv((G.T.dot(W).dot(G)))
print beta, wls_V

wls_fit = sm.WLS(Y, G, np.diag(W)).fit()
print wls_fit.params, wls_fit.cov_params()

beta = np.einsum('i...,ij...->...i', G, W_ill)
beta = 1. / np.einsum('...i,i...->...', beta, G)
beta = np.einsum('...,i...->...i', beta, G)
beta = np.einsum('...i,ij...->j...', beta, W_ill)
beta = np.einsum('i...,i...->...', beta, Z)

residuals = (Z - beta)
s2 = np.einsum('i...,ij...->...i', residuals, W_ill)
s2 = np.einsum('...i,i...->...', s2, residuals)
s2 /= (residuals.shape[0] - 1)

wls_V = 1. / W_ill.sum(0).sum(0) # inv(G.T.dot(W).dot(G))
wls_V *= s2

t_in_limbo = beta / np.sqrt(wls_V)
plt.imshow(t_in_limbo, cmap=plt.cm.hot)
plt.colorbar()

plt.title('t-values for in limbo effect')


threshold = sp.stats.t(len(ms_data) - 1).ppf(0.95)

plt.figure(figsize=(16,6))

cmap = plt.cm.hot
cmap.set_bad('black')

plt.subplot(131)
plt.imshow(z, cmap=cmap, vmin=z.min(), vmax=z.max())
plt.grid(False)
plt.title('Unthresholded z-map', fontsize=20)

plt.subplot(132)
plt.imshow(np.ma.masked_array(z_threshold, z_threshold == 0), cmap=cmap, vmin=z.min(), vmax=z.max())
plt.grid(False)
plt.title('Thresholded z-map', fontsize=20)

plt.subplot(133)
plt.imshow(np.ma.masked_where((z_threshold != 0) | ((z_threshold == 0) & (t_in_limbo > threshold)), z), 
            cmap=cmap,
            vmin=z.min(), vmax=z.max()           )

plt.title('In limbo regions', fontsize=20)

plt.grid(False)