import numpy as np
import plotly.graph_objs as go
from plotly.offline import download_plotlyjs, init_notebook_mode,  iplot, plot
from IPython.core.display import display, HTML
import scipy.io as sio
import os
# Comment out if locally 
os.chdir("/home/jovyan/content/04/subsection/02")

init_notebook_mode(connected=True)
config={'showLink': False, 'displayModeBar': False}

matrix = sio.loadmat(file_name = "newmap.mat", matlab_compatible = True)
newmap = matrix['vect']

_,r,c = newmap.shape

# Define frames
nb_frames = 42
matrix = np.rot90(newmap)

fig = go.Figure(frames=[go.Frame(data=go.Surface(
    z=(41 - k * 1) * np.ones((r, c)),
    surfacecolor=np.rot90(matrix[k,:,:], 2),
    cmin=0, cmax=80
    ),
    name=str(k) # you need to name the frame for the animation to behave properly
    )
    for k in range(nb_frames)])

# Add data to be displayed before animation starts
fig.add_trace(go.Surface(
    z=41 * np.ones((r, c)),
    surfacecolor=np.flipud(newmap[41]),
    colorscale='Viridis',
    cmin=0, cmax=80,
    colorbar=dict(thickness=20, ticklen=4)
    ))


def frame_args(duration):
    return {
            "frame": {"duration": duration},
            "mode": "immediate",
            "fromcurrent": True,
            "transition": {"duration": duration, "easing": "linear"},
        }

sliders = [
            {
                "pad": {"b": 10, "t": 60},
                "len": 0.9,
                "x": 0.1,
                "y": 0,
                "steps": [
                    {
                        "args": [[f.name], frame_args(0)],
                        "label": str(abs(42 - k)),
                        "method": "animate",
                    }
                    for k, f in enumerate(fig.frames)
                ],
            }
        ]

# Layout
fig.update_layout(
         title='Axial view slices',
         width=700,
         height=800,
         scene=dict(
                    zaxis=dict(range=[0, 41], autorange=False),
                    aspectratio=dict(x=0.9, y=0.9, z=0.9),
                    ),
         updatemenus = [
            {
                "buttons": [
                    {
                        "args": [None, frame_args(50)],
                        "label": "▶", # play symbol
                        "method": "animate",
                    },
                    {
                        "args": [[None], frame_args(0)],
                        "label": "◼", # pause symbol
                        "method": "animate",
                    },
                ],
                "direction": "left",
                "pad": {"r": 10, "t": 70},
                "type": "buttons",
                "x": 0.1,
                "y": 0,
            }
         ],
         sliders=sliders
)

# Create 3d animation figure 
plot(fig, filename = 'fig_3d.html', config = config)

# THEBELAB
display(HTML('fig_3d.html'))
# BINDER
# iplot(fig,config=config)