# Import data
import time
import numpy as np
from skimage import io
import plotly.graph_objects as go
from scipy.ndimage import uniform_filter1d
from scipy.signal import savgol_filter
from scipy import *
from scipy.signal import *
from util import *
"""
vol = io.imread("https://s3.amazonaws.com/assets.datacamp.com/blog_assets/attention-mri.tif")
volume = vol.T
r, c = volume[0].shape
"""

# Import data
dossier = "Dataset/Shear_Wave_Rot00_CSV_Data"
fichiers_selectionnes= fichiers_selectionnes = ['Shear_x001-x101_y{:03d}_Rot00.csv'.format(i) for i in range(10, 62)]
print(fichiers_selectionnes)

pre_volume= np.array(lire_fichier_csv(dossier,fichiers_selectionnes))
pre_traitement = abs(hilbert(pre_volume))

traitement = savgol_filter(pre_traitement, window_length=50, polyorder=5)

volume = traitement[:, ::10, :]
#volume = np.load('Dataset/3D_Dataset_All_Wave_PY_Data/3D_Dataset_shear_Wave_Rot00.npy')
r, c = volume[0].shape
# Define the number of frames based on the depth of the 3D volume
nb_frames = volume.shape[0]

# Define the z-scale based on the number of frames
z_scale = 6.7 / nb_frames

#custom_colorscale = [(0, 'blue'), (0.5, 'cyan'), (1, 'red')]

fig = go.Figure(frames=[go.Frame(data=go.Surface(
    z=(z_scale * nb_frames - k * z_scale) * np.ones((r, c)),
    surfacecolor=np.flipud(volume[nb_frames - 1 - k]),
    cmin=0, cmax=20000,
    colorscale='Jet'
    ),
    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=z_scale * nb_frames * np.ones((r, c)),
    surfacecolor=np.flipud(volume[nb_frames - 1]),
    colorscale='Jet',
    cmin=-20000, cmax=20000,
    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(k),
                        "method": "animate",
                    }
                    for k, f in enumerate(fig.frames)
                ],
            }
        ]

# Layout
fig.update_layout(
         title='Slices in volumetric data',
         width=600,
         height=600,
         scene=dict(
                    zaxis=dict(range=[-0.1, z_scale * nb_frames], autorange=False),
                    aspectratio=dict(x=1, y=1, z=1),
                    ),
         updatemenus = [
            {
                "buttons": [
                    {
                        "args": [None, frame_args(50)],
                        "label": "&#9654;", # play symbol
                        "method": "animate",
                    },
                    {
                        "args": [[None], frame_args(0)],
                        "label": "&#9724;", # pause symbol
                        "method": "animate",
                    },
                ],
                "direction": "left",
                "pad": {"r": 10, "t": 70},
                "type": "buttons",
                "x": 0.1,
                "y": 0,
            }
         ],
         sliders=sliders
)

fig.show()