import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyvista as pv
import pvxarray
import xarray as xr

pv.set_plot_theme("document")
pv.set_jupyter_backend("server")

# Temporary manual mode override; waiting for pyvista release including changes from https://github.com/pyvista/pyvista/pull/5798

pv.global_theme.trame.server_proxy_enabled = False
pv.global_theme.trame.jupyter_extension_enabled = True

ds = xr.open_dataset("data/all_flashes.nc")
ds

# event_parent_flash_id , flash_event_count (min 10)
ds.flash_id[0]

_ = plt.hist(ds.event_z, bins=100, range=(400, 15000))

_ = plt.hist(ds.event_parent_flash_id, bins=100)

ds.event_parent_flash_id[ds.event_parent_flash_id == 2000]

points = np.c_[ds.event_x, ds.event_y, ds.event_z]

t = pd.to_datetime(ds.event_time).astype(np.int64)


def get_flash(flash_index):
    flash_id = ds.event_parent_flash_id[flash_index]
    event_mask = ds.event_parent_flash_id == flash_id

    pc = pv.PolyData(points[event_mask])
    pc["event_time"] = t[event_mask]  # da.data
    pc["event_parent_flash_id"] = ds.event_parent_flash_id[event_mask]

    pc = pc.elevation().threshold((400, 15000), scalars="Elevation").extract_geometry()
    return pc


pc = pv.PolyData(points)
pc["event_time"] = t.values
pc["event_parent_flash_id"] = ds.event_parent_flash_id.values

full_pc = pc.elevation().threshold((400, 15000), scalars="Elevation")
box = full_pc.outline()

# full_pc.extract_surface().save('lightning.vtp')

pc = get_flash(2000)

pl = pv.Plotter()
actor = pl.add_mesh(
    pc,
    scalars="event_time",
    cmap="plasma",
    point_size=10,
    ambient=0.5,
)
pl.set_background("grey")
pl.show()

import ipywidgets as widgets
import pyvista as pv
from tqdm import tqdm

from pvxarray.vtk_source import PyVistaXarraySource


def time_controls(plotter: pv.BasePlotter, continuous_update=False, step=1):
    def update_time_index(time_index):
        ...
        plotter.add_mesh(
            get_flash(time_index),
            scalars="event_time",
            cmap="plasma",
            point_size=10,
            ambient=0.5,
            render_points_as_spheres=True,
            name="flash",
        )

        plotter.render()

    tmax = len(ds.event_time)

    def set_time(change):
        value = change["new"]
        if value < 0:
            value = 0
        if value >= tmax:
            value = tmax - 1
        update_time_index(value)

    play = widgets.Play(
        value=0,
        min=0,
        max=tmax,
        step=step,
        description="Time Index",
    )
    play.observe(set_time, "value")

    slider = widgets.IntSlider(min=0, max=tmax, continuous_update=continuous_update)
    widgets.jslink((play, "value"), (slider, "value"))
    return widgets.HBox([play, slider])


pl = pv.Plotter()
actor = pl.add_mesh(
    get_flash(0),
    scalars="event_time",
    cmap="plasma",
    point_size=10,
    ambient=0.5,
    name="flash",
    render_points_as_spheres=True,
)
pl.add_mesh(box, color="k")
pl.set_background("grey")
pl.show()

time_controls(pl, continuous_update=True)