import streamlit as st
import plotly.graph_objects as go
import numpy as np

def plot(kvl_file):
    st.title('KVL Data Visualisation')
    nt = 0
    sampleno = 0

    if not kvl_file:
        st.warning('No file has been uploaded')
    else:
        if 'TIME' in kvl_file:
            nt = len(kvl_file['TIME'])
            dt = int(kvl_file['TIME'][1]) - int(kvl_file['TIME'][0])
            t_iter = list(kvl_file['TIME'])
            keys = list(kvl_file.keys())

            curves = st.multiselect('Select TOP Key to Visualize surface, SED keys to project', keys)
            inst = st.select_slider('Select Time Instance', t_iter)
            gridno=st.checkbox('Grid on')
            if gridno:
                kgrid=st.number_input('Higher the number, coarser the grid',value=2,placeholder=2)
            else: 
                kgrid=2

            time = int(inst) / dt
            sampleno = len(kvl_file['TOP']) / nt
            start_index = int(time * sampleno)
            end_index = start_index + int(sampleno)

            # Get the data for the current instance
            for curve in curves:
                height_data = kvl_file['TOP'][start_index:end_index]
                intensity_data= kvl_file[curve][start_index:end_index]


                # Create the surface plot for the 'TOP' key
                z_data = np.array(height_data)
                x_data = np.arange(z_data.shape[1])
                y_data = np.arange(z_data.shape[0])
                
                # Create meshgrid for surface
                xGrid, yGrid = np.meshgrid(x_data[::kgrid], y_data[::kgrid])
                xRev, yRev = np.meshgrid(y_data[::kgrid], x_data[::kgrid])

                # Create figure
                fig = go.Figure(data=[go.Surface(x=x_data, y=y_data, z=z_data,surfacecolor=intensity_data,colorscale='Rainbow')])
                if gridno: #grid on 
                    for i,j,k in zip(xGrid,yGrid,z_data[::kgrid,::kgrid]):
                        fig.add_trace(go.Scatter3d(x=i,y=j,z=k, mode='lines', line=dict(color='black', width=1,)))
                    for i,j,k in zip(xRev,yRev,z_data[::kgrid,::kgrid].T):
                        fig.add_trace(go.Scatter3d(x=j,y=i,z=k, mode='lines', line=dict(color='black', width=1,)))
        
                # Update the layout of the figure
                fig.update_layout(
                    height=700,
                    showlegend=False,
                    title='KVL Data Visualisation of '+curve,
                    scene=dict(
                        xaxis_title='X',
                        yaxis_title='Y',
                        zaxis_title='Z',
                        aspectmode='manual',
                        aspectratio=dict(x=1.5, y=1.5, z=1),
                        camera=dict(
                            eye=dict(x=2.2, y=-1, z=0.6)
                        ),
                    )
                )

                # Display the figure
                st.plotly_chart(fig, theme=None, use_container_width=True)

        else:
            st.warning('No time data found')