import os
import tempfile

from IPython.display import Image, display
import nibabel as ni

import osmosis.viz.maya as viz
import osmosis.utils as ozu
import osmosis.model.dti as dti
import osmosis.model.analysis as oza


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

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')

vol_anat = oio.get_t1(subject, resample=ni.load(oio.data_path + '/%s/%s_wm_mask.nii.gz'%(subject, subject)))

rrmse_data = [oza.cross_predict(TM_1k_1, TM_1k_2), oza.cross_predict(TM_2k_1, TM_2k_2), oza.cross_predict(TM_4k_1, TM_4k_2)]

%gui wx
fn = []
for vol_rmse in rrmse_data:
    fn.append('%s.png'%tempfile.NamedTemporaryFile().name)
    viz.plot_cut_planes(vol_anat,
                        overlay=vol_rmse,
                        vmin=0.5,
                        vmax=1.5,
                        overlay_cmap="RdYlGn",
                        invert_cmap=True,
                        slice_coronal=40,
                        slice_saggital=15,
                        slice_axial=30,
                        view_azim=-40,
                        view_elev=60,
                        file_name=fn[-1])

for this_fn in fn:
    i = Image(filename=this_fn, width=1280, height=1024)
    display(i)

fn = []
for vol_rmse in rrmse_data:
    fn.append('%s.png'%tempfile.NamedTemporaryFile().name)
    viz.plot_cut_planes(vol_anat,
                        overlay=vol_rmse,
                        vmin=0.5,
                        vmax=1.5,
                        overlay_cmap="RdYlGn",
                        invert_cmap=True,
                        slice_coronal=40,
                        slice_saggital=15,
                        slice_axial=45,
                        view_azim=40,
                        view_elev=60,
                        file_name=fn[-1])

for this_fn in fn:
    i = Image(filename=this_fn, width=1280, height=1024)
    display(i)

viz.plot_cut_planes(vol_anat,
                    overlay=rrmse_data[-1],
                    vmin=0.5,
                    vmax=1.5,
                    overlay_cmap="RdYlGn",
                    invert_cmap=True,
                    slice_coronal=None,
                    slice_saggital=None,
                    slice_axial=45,
                    view_azim=0,
                    view_elev=0)

%gui wx
viz.plot_cut_planes(vol_anat,
                    overlay=rrmse_data[-1],
                    vmin=0.5,
                    vmax=1.5,
                    overlay_cmap="RdYlGn",
                    invert_cmap=True,
                    slice_coronal=None,
                    slice_saggital=None,
                    slice_axial=30,
                    view_azim=0,
                    view_elev=0)