# Prepare Python environment

import scipy.io as sio
from pathlib import Path

data_dir = Path("../../../data/03-T2")
data_file = "T2w_vs_T2map.mat"

#Load either archived or generated plot variables
mat_contents = sio.loadmat(data_dir / data_file)

# Get MATLAB data
T2_map = mat_contents['T2_map']
TE_1 = mat_contents['TE_1']
TE_2 = mat_contents['TE_2']
TE_3 = mat_contents['TE_3']
TE_4 = mat_contents['TE_4']

# Crop the images so that the brain is centered on the figure
T2_map = T2_map[:-100, :]
TE_1 = TE_1[:-100, :]
TE_2 = TE_2[:-100, :]
TE_3 = TE_3[:-100, :]
TE_4 = TE_4[:-100, :]

## Plot

import matplotlib.pyplot as plt
import chart_studio.plotly as py
import plotly.graph_objs as go
import numpy as np
from chart_studio.plotly import plot
from IPython.display import display, HTML
from plotly import tools
from plotly.offline import download_plotlyjs, init_notebook_mode, plot, iplot

config={'showLink': False, 'displayModeBar': False}

xAxis = np.arange(0, T2_map.shape[1])
yAxis = np.arange(0, T2_map.shape[0])

trace1 = go.Heatmap(x = xAxis,
                   y = yAxis,
                   z=np.flipud(TE_1),
                   colorscale='Gray',
                   showscale = False,
                   visible=False,
                   name = 'Signal')
trace2 = go.Heatmap(x = xAxis,
                   y = yAxis,
                   z=np.flipud(TE_2),
                   colorscale='Gray',
                   showscale = False,
                   visible=False,
                   name = 'Signal')
trace3 = go.Heatmap(x = xAxis,
                   y = yAxis,
                   z=np.flipud(TE_3),
                   colorscale='Gray',
                   showscale = False,
                   visible=True,
                   name = 'Signal')
trace4 = go.Heatmap(x = xAxis,
                   y = yAxis,
                   z=np.flipud(TE_4),
                   colorscale='Gray',
                   visible=False,
                   showscale = False,
                   name = 'Signal')
trace5 = go.Heatmap(x = xAxis,
                   y = yAxis,
                   z=np.flipud(T2_map),
                   colorscale='Portland',
                   xaxis='x2',
                   yaxis='y2',
                   visible=True,
                   name = 'T1 values (ms)',
                   zmin=100, 
                   zmax=200,)


data=[trace1, trace2, trace3, trace4, trace5]

updatemenus = list([
    dict(active=2,
         x = 0.12,
         xanchor = 'left',
         y = -0.15,
         yanchor = 'bottom',
         direction = 'up',
         font=dict(
                family='Times New Roman',
                size=16
            ),
         buttons=list([   
            dict(label = '12.8 ms',
                 method = 'update',
                 args = [{'visible': [True, False, False, False, True]},
                         ]),
            dict(label = '128 ms',
                 method = 'update',
                 args = [{'visible': [False, True, False, False, True]},
                         ]),
            dict(label = '256 ms',
                 method = 'update',
                 args = [{'visible': [False, False, True, False, True]},
                         ]),
            dict(label = '384 ms',
                 method = 'update',
                 args = [{'visible': [False,False, False, True, True]},
                         ])
        ]),
    )
])

layout = dict(
    width=728,
    height=448.5,
    margin = dict(
                t=90,
                r=140,
                b=90,
                l=30),
    annotations=[
        dict(
            x=0.13,
            y=1.15,
            showarrow=False,
            text='T2w images',
            font=dict(
                family='Times New Roman',
                size=26
            ),
            xref='paper',
            yref='paper'
        ),
        dict(
            x=0.82,
            y=1.15,
            showarrow=False,
            text='T<sub>2</sub> map',
            font=dict(
                family='Times New Roman',
                size=26
            ),
            xref='paper',
            yref='paper'
        ),
        dict(
            x=1.17,
            y=1.15,
            showarrow=False,
            text='T<sub>2</sub> (ms)',
            font=dict(
                family='Times New Roman',
                size=26
            ),
            xref='paper',
            yref='paper'
        ),
        dict(
            x=0.02,
            y=-0.15,
            showarrow=False,
            text='TE:',
            font=dict(
                family='Times New Roman',
                size=22
            ),
            xref='paper',
            yref='paper'
        ),
    ],
    xaxis=dict(range=[0, xAxis[-1]], autorange=False,
               showgrid=False, zeroline=False, showticklabels=False,
               ticks='', domain=[0, 0.5]),  # First plot domain
    yaxis=dict(range=[0, yAxis[-1]], autorange=False,
               showgrid=False, zeroline=False, showticklabels=False,
               ticks='', domain=[0, 1]),
    xaxis2=dict(range=[0, xAxis[-1]], autorange=False,
                showgrid=False, zeroline=False, showticklabels=False,
                ticks='', domain=[0.5, 1]),  # Second plot domain
    yaxis2=dict(range=[0, yAxis[-1]], autorange=False,
                showgrid=False, zeroline=False, showticklabels=False,
                ticks='', domain=[0, 1], anchor='x2'),
    showlegend=False,
    autosize=False,
    updatemenus=updatemenus
)

fig = dict(data=data, layout=layout)

iplot(fig, filename = 'ir_fig_2.html', config = config)