import os
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
    # -----------------------------------------------------------------------------
    ds = xr.open_dataset(os.path.join(CURRENT_DIRECTORY, "../../../../../Data/Field_20251119_140509_BuildingGroupFloor.nc"))

    # ----- Load fluid field -----
    # 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, :, :]
    nTime, nz, ny, nx, _ = velFieldRaw.shape
    print(f"Fluid field resolution: {nx=}, {ny=}, {nz=}")

    # ----- 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=}")

    # ----- Fluid field at the first stored time step -----
    velField0 = velFieldRaw[0, :, :, :, :3]
    velFieldMag0 = np.linalg.norm(velField0.values, axis=-1).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)
    # -----------------------------------------------------------------------------

    modelReader = vtk.vtkPLYReader()
    # modelReader.SetFileName("Data/SmallBuilding3.ply")
    modelReader.SetFileName(os.path.join(CURRENT_DIRECTORY, "../../../../Data/Buildings-GroupAll.ply"))
    modelReader.Update()

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

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

    transform = vtk.vtkTransform()

    transform.PostMultiply()

    # Small building 3
    # transform.RotateX(-90.0)
    # transform.RotateY(50.0)

    # Building Group
    transform.Translate(25, 0.2, -12)
    transform.RotateX(180.0)
    transform.RotateY(180.0)

    modelActor.SetUserTransform(transform)

    # -----------------------------------------------------------------------------
    # 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.ResetCamera()
    renderWindow.Render()

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

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

    @state.change("time")
    def update_time(time, **kwargs):
        time_index = round(time / dt) - 1
        newVelField = velFieldRaw[time_index, :, :, :, :3]
        newVelMagFlat = np.linalg.norm(newVelField.values, axis=-1).ravel()
        grid.point_data["vel_mag"] = newVelMagFlat
        grid.Modified()
        ctrl.view_update()

    @state.change("ypos")
    def update_ypos(ypos, **kwargs):
        plane.SetOrigin(center_x, ypos, center_z)
        ctrl.view_update()

    def reset_ypos():
        state.ypos = center_y

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

        with layout.toolbar:
            vuetify.VSpacer()

            # 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;",
            )
            with vuetify.VBtn(icon=True, click=reset_ypos):
                vuetify.VIcon("mdi-restore")
            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()

if __name__ == "__main__":
    main()