import os
import asyncio

from trame.app import get_server
from trame.ui.vuetify import SinglePageLayout
from trame.widgets import vuetify
from trame.widgets import vtk as vtk_widget
import vtk
import xarray as xr
import pyvista as pv
import numpy as np

import matplotlib.pyplot as plt
from vtkmodules.vtkCommonCore import vtkLookupTable

from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonDataModel import vtkPlane
from vtkmodules.vtkFiltersCore import vtkCutter
from vtkmodules.vtkFiltersModeling import vtkOutlineFilter
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderer,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
)

def main(server=None, **kwargs):
    CURRENT_DIRECTORY = os.path.abspath(os.path.dirname(__file__))

    # -----------------------------------------------------------------------------
    # NetCDF to VTK
    # -----------------------------------------------------------------------------
    nc_path_1 = os.path.join(CURRENT_DIRECTORY, "data/Field_20251122_225911_BuildingGroupFloor.nc")
    nc_path_2 = os.path.join(CURRENT_DIRECTORY, "data/Field_20251122_224941_BuildingGroupFloor_Moved.nc")

    ds = xr.open_dataset(nc_path_1)
    ds_moved = xr.open_dataset(nc_path_2)

    # ----- Load fluid field 1 -----
    # Vtk uses right-handed coordinate system, while unity uses left-handed coordinate system. So the 
    # z asix needs to be reversed here.
    velFieldRaw = ds["velField"][:, ::-1, :, :, :3]
    velFieldMag = np.linalg.norm(velFieldRaw.values, axis=-1)
    nTime, nz, ny, nx, _ = velFieldRaw.shape
    print(f"Fluid field resolution: {nx=}, {ny=}, {nz=}")

    # ----- Load fluid field 2 (moved) -----
    velFieldMovedRaw = ds_moved["velField"][:, ::-1, :, :, :3]
    velFieldMovedMag = np.linalg.norm(velFieldMovedRaw.values, axis=-1)

    # ----- Load dx -----
    velXCoor = ds["velX"]
    dx = float(velXCoor.values[1] - velXCoor.values[0])
    print(f"{dx=}")

    # ----- Load dt -----
    timeList = ds["time"]
    dt = float(timeList.values[1] - timeList.values[0])
    dt = round(dt, 2)
    print(f"{dt=}")

    # ----- Initial fluid field -----
    velFieldMag0 = velFieldMag[0, :, :, :].ravel()

    # ----- Build vtk data -----
    grid = pv.ImageData()
    grid.dimensions = (nx, ny, nz)
    grid.origin = (-(nx-1) * dx / 2, -(ny-1) * dx / 2 + 10, -(nz-1) * dx / 2)
    grid.spacing = (dx, dx, dx)
    grid.point_data["vel_mag"] = velFieldMag0

    vtk_producer = vtk.vtkTrivialProducer()
    vtk_producer.SetOutput(grid)

    # -----------------------------------------------------------------------------
    # PLY Reader (for building model)
    # -----------------------------------------------------------------------------

    # ----- Transform logic -----
    transform = vtk.vtkTransform()
    transform.PostMultiply()
    transform.Translate(25, 0.2, -12)
    transform.RotateX(180.0)
    transform.RotateY(180.0)

    # ----- Obstacle for fluid field 1 -----
    ply_path_1 = os.path.join(CURRENT_DIRECTORY, "data/Buildings-GroupAll.ply")
    modelReader = vtk.vtkPLYReader()
    modelReader.SetFileName(ply_path_1)
    modelReader.Update()

    modelMapper = vtk.vtkPolyDataMapper()
    modelMapper.SetInputConnection(modelReader.GetOutputPort())

    modelActor = vtk.vtkActor()
    modelActor.SetMapper(modelMapper)
    modelActor.SetUserTransform(transform)

    # ----- Obstacle for fluid field 2 -----
    ply_path_2 = os.path.join(CURRENT_DIRECTORY, "data/Buildings-GroupAll-Moved.ply")
    modelReader2 = vtk.vtkPLYReader()
    modelReader2.SetFileName(ply_path_2)
    modelReader2.Update()

    modelMapper2 = vtk.vtkPolyDataMapper()
    modelMapper2.SetInputConnection(modelReader2.GetOutputPort())

    modelActor2 = vtk.vtkActor()
    modelActor2.SetMapper(modelMapper2)
    modelActor2.SetUserTransform(transform)
    modelActor2.SetVisibility(False)

    # -----------------------------------------------------------------------------
    # Render VTK fluid field
    # -----------------------------------------------------------------------------

    renderer = vtkRenderer()
    renderWindow = vtkRenderWindow()
    renderWindow.AddRenderer(renderer)

    renderWindowInteractor = vtkRenderWindowInteractor()
    renderWindowInteractor.SetRenderWindow(renderWindow)
    renderWindowInteractor.GetInteractorStyle().SetCurrentStyleToTrackballCamera()

    # ----- Render outline -----

    colors = vtkNamedColors()

    outline = vtkOutlineFilter()
    outline.SetInputConnection(vtk_producer.GetOutputPort())

    outlineMapper = vtkPolyDataMapper()
    outlineMapper.SetInputConnection(outline.GetOutputPort())

    outlineActor = vtkActor()
    outlineActor.SetMapper(outlineMapper)
    outlineActor.GetProperty().SetColor(colors.GetColor3d("White"))

    # ----- Colormap -----
    def create_vtk_lut_from_matplotlib(cmap_name, num_colors=256):
        mpl_cmap = plt.get_cmap(cmap_name)
        
        lut = vtkLookupTable()
        lut.SetNumberOfTableValues(num_colors)
        
        for i in range(num_colors):
            rgba = mpl_cmap(i / (num_colors - 1))
            lut.SetTableValue(i, rgba[0], rgba[1], rgba[2], 1.0)
        
        lut.SetRange(0, 15) 
        lut.Build()
        return lut

    turbo_lut = create_vtk_lut_from_matplotlib("turbo")

    # ----- Render a slice -----
    center_x = grid.center[0]
    center_y = grid.center[1]
    center_z = grid.center[2]
    plane = vtkPlane()
    plane.SetOrigin(center_x, center_y, center_z)
    plane.SetNormal(0, 1, 0)
    print(f"{grid.center[1]=}")

    cutter = vtkCutter()
    cutter.SetInputConnection(vtk_producer.GetOutputPort())
    cutter.SetCutFunction(plane)
    cutter.Update()

    slice_mapper = vtkPolyDataMapper()
    slice_mapper.SetInputConnection(cutter.GetOutputPort())
    slice_mapper.SetScalarModeToUsePointFieldData()
    slice_mapper.SelectColorArray("vel_mag")
    slice_mapper.SetLookupTable(turbo_lut)
    slice_mapper.SetScalarRange(0, 15)

    slice_actor = vtkActor()
    slice_actor.SetMapper(slice_mapper)

    # Add the actors to the renderer
    renderer.AddActor(outlineActor)
    renderer.AddActor(slice_actor)
    renderer.AddActor(modelActor)
    renderer.AddActor(modelActor2)
    renderer.ResetCamera()
    renderWindow.Render()

    # -----------------------------------------------------------------------------
    # GUI
    # -----------------------------------------------------------------------------

    server = get_server(client_type = "vue2")
    ctrl = server.controller
    state = server.state

    # state.trame__title = "Visualizer"
    state.time = dt
    state.ypos = center_y
    state.show_moved_case = False
    state.show_diff = False
    state.is_playing = False

    # --- Animation Logic ---
    async def play_loop():
        """Asynchronous loop to update time"""
        while state.is_playing:
            with state:
                new_time = state.time + dt
                
                if new_time > nTime * dt:
                    new_time = dt
                
                state.time = new_time
            
            await asyncio.sleep(0.75)
    
    def toggle_play():
        """Toggle the animation state"""
        state.is_playing = not state.is_playing
        if state.is_playing:
            asyncio.create_task(play_loop())

    @state.change("time")
    def update_time(time, **kwargs):
        time_index = round(time / dt) - 1

        if state.show_diff:
            new_vel_field = velFieldRaw[time_index, :, :, :, :].values
            new_vel_field_moved = velFieldMovedRaw[time_index, :, :, :, :].values
            new_vel_field_diff = new_vel_field_moved - new_vel_field
            new_scalar_field = np.linalg.norm(new_vel_field_diff, axis=-1).ravel()
        elif state.show_moved_case:
            new_scalar_field = velFieldMovedMag[time_index, :, :, :].ravel()
        else:
            new_scalar_field = velFieldMag[time_index, :, :, :].ravel()

        grid.point_data["vel_mag"] = new_scalar_field
        grid.Modified()
        ctrl.view_update()

    @state.change("ypos")
    def update_ypos(ypos, **kwargs):
        plane.SetOrigin(center_x, ypos, center_z)
        ctrl.view_update()
    
    @state.change("show_moved_case")
    def toggle_case(show_moved_case, **kwargs):
        """Switch visibility of obstacles and refresh fluid data"""
        modelActor.SetVisibility(not show_moved_case)
        modelActor2.SetVisibility(show_moved_case)
        update_time(state.time)
        ctrl.view_update()
    
    @state.change("show_diff")
    def toggle_diff(show_diff, **kwargs):
        """Toggle difference mode"""
        update_time(state.time)
        ctrl.view_update()

    with SinglePageLayout(server) as layout:
        layout.title.set_text("Fluid Field Visualizer")

        with layout.toolbar:
            vuetify.VSpacer()

            # --- Switch 1: original or moved ---
            vuetify.VSwitch(
                v_model=("show_moved_case", False),
                label="Show field (moved)",
                hide_details=True,
                dense=True,
                classes="mr-4 mt-4",
                style="max-width: 200px;",
            )

            vuetify.VDivider(vertical=True, classes="mx-2")

            # --- Switch 2: normal or difference ---
            vuetify.VSwitch(
                v_model=("show_diff", False),
                label="Show field difference",
                color="red", # Highlight this switch
                hide_details=True,
                dense=True,
                classes="mr-4 mt-4",
                style="max-width: 200px;",
            )

            vuetify.VDivider(vertical=True, classes="mx-2")

            # --- Play/Pause Button ---
            with vuetify.VBtn(label="Play", icon=True, click=toggle_play, classes="mr-4 mt-4"):
                vuetify.VIcon("{{ is_playing ? 'mdi-pause' : 'mdi-play' }}")

            # Time slider
            vuetify.VSlider(
                label="Time (s)",
                v_model=("time", 0),
                min=dt,
                max=nTime * dt,
                step=dt,
                thumb_label=True,
                hide_details=True,
                dense=True,
                style="max-width: 300px; margin-top: 40px; margin-bottom: 10px;",
            )

            vuetify.VDivider(vertical=True, classes="mx-2")

            # Slice position silder
            vuetify.VSlider(
                label="Y Position (m)",
                v_model=("ypos", center_y),
                min=grid.bounds[2] + dx, 
                max=grid.bounds[3] - dx, 
                step=dx,
                thumb_label=True,
                hide_details=True, 
                dense=True,
                style="max-width: 300px; margin-top: 40px; margin-bottom: 10px;",
            )
            vuetify.VDivider(vertical=True, classes="mx-2")

        with layout.content:
            with vuetify.VContainer(
                fluid=True,
                classes="pa-0 fill-height",
            ):
                view = vtk_widget.VtkLocalView(renderWindow)
                ctrl.view_update = view.update

    server.start(timeout=0)

if __name__ == "__main__":
    main()