%pylab inline

import nibabel as ni
import osmosis.model.analysis as oza
import osmosis.model.sparse_deconvolution as sfm
import osmosis.model.dti as dti
import osmosis.viz.mpl as mpl
import osmosis.boot as boot
import osmosis.utils as ozu

n_bvecs= [6, 10, 20, 30, 40, 80, 120]
#n_bvecs= [6, 10, 20, 40]
n_reps = 1
import matplotlib.colors as co
cNorm  = co.Normalize(vmin=0, vmax=len(n_bvecs)+1)
scalarMap = cm.ScalarMappable(norm=cNorm, cmap=matplotlib.cm.rainbow)
colors = scalarMap.to_rgba(range(len(n_bvecs)+1))

import os
import osmosis as oz
import osmosis.io as oio
oio.data_path = os.path.join(oz.__path__[0], 'data')

bvals = [1000, 2000, 4000]
subject = 'SUB1'
data_1k_1, data_1k_2 = oio.get_dwi_data(1000, subject)
data_2k_1, data_2k_2 = oio.get_dwi_data(2000, subject)
data_4k_1, data_4k_2 = oio.get_dwi_data(4000, subject)

wm_mask = np.zeros(ni.load(data_1k_1[0]).shape[:3])
wm_nifti = ni.load(oio.data_path + '/%s/%s_wm_mask.nii.gz'%(subject, subject)).get_data()
wm_mask[np.where(wm_nifti==1)] = 1

TM_1k_1 = dti.TensorModel(*data_1k_1, mask=wm_mask, params_file='temp')
TM_1k_2 = dti.TensorModel(*data_1k_2, mask=wm_mask, params_file='temp')
TM_2k_1 = dti.TensorModel(*data_2k_1, mask=wm_mask, params_file='temp')
TM_2k_2 = dti.TensorModel(*data_2k_2, mask=wm_mask, params_file='temp')
TM_4k_1 = dti.TensorModel(*data_4k_1, mask=wm_mask, params_file='temp')
TM_4k_2 = dti.TensorModel(*data_4k_2, mask=wm_mask, params_file='temp')

rrmse = np.empty((3, len(n_bvecs), n_reps) + wm_mask.shape,)
rrmse_150 = np.empty((3, ) + wm_mask.shape)
rrmse.shape, rrmse_150.shape

fig_curves, ax_curves = plt.subplots(1, 3)
fig_curves.set_size_inches([12, 6])

for Model in [dti.TensorModel, sfm.SparseDeconvolutionModel]:
    for b_idx, (M1, M2) in enumerate(zip([TM_1k_1, TM_2k_1, TM_4k_1], [TM_1k_2, TM_2k_2, TM_4k_2])):
        M1 = Model(M1.data, M1.bvecs, M1.bvals * 1000, mask=wm_mask, params_file='temp')
        M2 = Model(M2.data, M2.bvecs, M2.bvals * 1000, mask=wm_mask, params_file='temp')
        rrmse_150[b_idx]= oza.cross_predict(M1, M2)
        
        for n_idx, n in enumerate(n_bvecs): 
            for rep in range(n_reps):
                new_bvecs, new_bv_i = boot.subsample(M1.bvecs[:, M1.b_idx], n)
                data_resamp_1 = np.concatenate([M1.data[..., M1.b0_idx], M1.data[..., M1.b_idx[new_bv_i]]], -1)
                bvecs_resamp_1 = np.concatenate([M1.bvecs[..., M1.b0_idx], M1.bvecs[..., M1.b_idx[new_bv_i]]], -1)
                bvals_resamp_1 = np.concatenate([M1.bvals[..., M1.b0_idx], M1.bvals[..., M1.b_idx[new_bv_i]]], -1)*1000
                data_resamp_2 = np.concatenate([M2.data[..., M2.b0_idx], M2.data[..., M2.b_idx[new_bv_i]]], -1)
                bvecs_resamp_2 = np.concatenate([M2.bvecs[..., M2.b0_idx], M2.bvecs[..., M2.b_idx[new_bv_i]]], -1)
                bvals_resamp_2 = np.concatenate([M2.bvals[..., M2.b0_idx], M2.bvals[..., M2.b_idx[new_bv_i]]], -1)*1000
                # Make sure you have the right number of regressors in the SFM case
                if Model==sfm.SparseDeconvolutionModel:
                    M1_resamp = Model(data_resamp_1, bvecs_resamp_1, bvals_resamp_1, mask=wm_mask, over_sample=150, params_file='temp')
                    M2_resamp = Model(data_resamp_2, bvecs_resamp_2, bvals_resamp_2, mask=wm_mask, over_sample=150, params_file='temp')
                else:
                    M1_resamp = Model(data_resamp_1, bvecs_resamp_1, bvals_resamp_1, mask=wm_mask, params_file='temp')
                    M2_resamp = Model(data_resamp_2, bvecs_resamp_2, bvals_resamp_2, mask=wm_mask, params_file='temp')

                rrmse[b_idx, n_idx, rep] = oza.cross_predict(M1_resamp, M2_resamp)
                
    
    median_rrmse = []
    median_rrmse_err = []

    for b_idx in range(3):
        median_rrmse.append([])
        median_rrmse_err.append([])
        fig = mpl.probability_hist(rrmse_150[b_idx][np.isfinite(rrmse_150[b_idx])], label='nbvecs=150', color=colors[-1])
        print("For b=%s, 150 dirs, the median PDD reliability is %2.2f"%(bvals[b_idx], np.median(rrmse_150[b_idx][np.isfinite(rrmse_150[b_idx])])))
        median_rrmse[-1] = [np.median(rrmse_150[b_idx][np.isfinite(rrmse_150[b_idx])])]
        median_rrmse_err[-1] = [0]
        for i in range(len(n_bvecs))[::-1]: 
            my_mean=(np.mean([np.median(rrmse[b_idx, i, j][np.isfinite(rrmse[b_idx, i, j])]) for j in range(n_reps)]))
            my_var=(np.var([np.median(rrmse[b_idx, i, j][np.isfinite(rrmse[b_idx, i, j])]) for j in range(n_reps)]))
            fig = mpl.probability_hist(rrmse[b_idx, i, j][np.isfinite(rrmse[b_idx, i, j])], fig=fig, label='nbvecs=%s'%str(n_bvecs[i]), color=colors[i])
            median_rrmse[-1].append(my_mean)
            median_rrmse_err[-1].append(my_var)
            print("For b=%s, %s dirs, the median PDD reliability is %2.2f"%(bvals[b_idx], n_bvecs[i], median_rrmse[-1][-1]))
            ax = fig.axes[0]

        ax.set_xlim([0.5,1.5])
        ax.set_xlabel('Relative RMSE')
        ax.set_ylabel('P(Relative RMSE)')
        plt.legend()
    
    for b_idx in range(3):
        ax_curves[b_idx].plot(np.hstack([150, n_bvecs[::-1]]), median_rrmse[b_idx], 'o-')
        #ax.errorbar(np.hstack([150, n_bvecs[::-1]]), median_rrmse[b_idx], median_rrmse_err[b_idx])
        ax_curves[b_idx].set_ylabel('Median Relative RMSE ')
        ax_curves[b_idx].set_xlabel('# b-vectors')
        ax_curves[b_idx].set_ylim([0.707, 1.05])
        ax_curves[b_idx].set_xticks(n_bvecs)
        ax_curves[b_idx].grid(True)
fig_curves.savefig('/home/arokem/Dropbox/curves_resample_plot.svg')