#!/usr/bin/env python
# coding: utf-8

# # Explore FVCOM GOM3 Hindcast on AWS Public Data - Interactive

# ## Original Imports and Setup
import xarray as xr
import numpy as np
import holoviews as hv
import geoviews as gv
import cartopy.crs as ccrs
import hvplot.xarray
import holoviews.operation.datashader as dshade
import pandas as pd
import datetime # Added for datetime object manipulation

# ## Panel Import and Extension
import panel as pn
pn.extension(sizing_mode="stretch_width") # Makes Panel objects responsive

# Configure HoloViews and Datashader
dshade.datashade.precompute = True
# hv.extension('bokeh') # pn.extension() handles this

# ## Load Dataset
url = 's3://umassd-fvcom/gom3/hindcast/parquet/combined.parq'
so = dict(anon=True)
print("Loading dataset (this might take a moment)...")
ds = xr.open_dataset(url, engine='kerchunk', chunks={'time':1},
                     backend_kwargs=dict(storage_options=dict(target_options=so,
                     remote_protocol='s3', lazy=True, remote_options=so)))
print("Dataset loaded.")
# print(ds) # Optional: display dataset summary

# ## Pre-calculate static components for the mesh
print("Pre-calculating mesh geometry...")
lon_vals = ds['lon'].load().data
lat_vals = ds['lat'].load().data
tris_df = pd.DataFrame(ds['nv'].T.values.astype('int') - 1, columns=['v0', 'v1', 'v2'])
tiles = gv.tile_sources.OSM
print("Mesh geometry pre-calculation complete.")

# ## Define Widgets
print("Setting up widgets...")

# Variable selection
plottable_variables = [
    var for var in ds.data_vars
    if 'time' in ds[var].dims and \
       'siglay' in ds[var].dims and \
       var not in ['lon', 'lat', 'lonc', 'latc', 'siglay', 'siglev', 'time', 'nv', 'nbe']
]
if not plottable_variables:
    raise ValueError("No suitable plottable variables found with 'time' and 'siglay' dimensions.")

var_widget = pn.widgets.Select(
    name='📈 Variable',
    options=plottable_variables,
    value=plottable_variables[0] if 'temp' not in plottable_variables else 'temp'
)

# Vertical level selection
num_siglay_levels = len(ds.siglay)
level_widget = pn.widgets.IntSlider(
    name='🌊 Vertical Level Index (0=surface)',
    start=0,
    end=num_siglay_levels - 1,
    step=1,
    value=0
)

# Time step selection using DatetimePicker
# Get the time range from the dataset
min_time_np = ds.time.min().values
max_time_np = ds.time.max().values

# Convert numpy.datetime64 to Python datetime.datetime for the widget
# Using pd.to_datetime handles the conversion robustly
min_time_dt = pd.to_datetime(str(min_time_np)).to_pydatetime()
max_time_dt = pd.to_datetime(str(max_time_np)).to_pydatetime()
# Set initial value, e.g., the first available time or a specific interesting time
initial_time_dt = pd.to_datetime(str(ds.time[0].values)).to_pydatetime()
# Fallback if a specific time was intended like in the original static example
# default_interesting_time_str = '2015-11-01 04:00:00'
# try:
# initial_time_dt = pd.to_datetime(default_interesting_time_str).to_pydatetime()
# if not (min_time_dt <= initial_time_dt <= max_time_dt):
# initial_time_dt = pd.to_datetime(str(ds.time[0].values)).to_pydatetime()
# except ValueError:
# initial_time_dt = pd.to_datetime(str(ds.time[0].values)).to_pydatetime()


time_widget = pn.widgets.DatetimePicker(
    name='⏰ Time Step (UTC)',
    start=min_time_dt,
    end=max_time_dt,
    value=initial_time_dt,
    # You can enable seconds if needed: enable_seconds=True,
    # step (for spinners, in seconds): step=3600 # e.g. 1 hour
)
print("Widgets setup complete.")

# ## Define the Plotting Function
@pn.depends(var_name=var_widget, level_idx=level_widget, selected_dt_obj=time_widget)
def create_fvcom_plot(var_name, level_idx, selected_dt_obj):
    print(f"Updating plot for: Variable={var_name}, Level Index={level_idx}, Time={selected_dt_obj.strftime('%Y-%m-%d %H:%M:%S')}")

    # Convert the datetime object from the picker to numpy.datetime64 for xarray
    # xarray's sel usually handles Python datetime objects well, but explicit conversion is safe.
    selected_time_np = np.datetime64(selected_dt_obj)

    print(f"Loading data: ds['{var_name}'].isel(siglay={level_idx}).sel(time='{selected_time_np}', method='nearest')")
    try:
        values_da = ds[var_name].isel(siglay=level_idx).sel(time=selected_time_np, method='nearest')
        values = values_da.load().data
        actual_time_selected = pd.to_datetime(str(values_da.time.data))
    except Exception as e:
        return pn.pane.Alert(f"Error loading data for {var_name}, level {level_idx}, time {selected_dt_obj}: {e}", alert_type='danger')

    print(f"Data loaded successfully. Min: {values.min()}, Max: {values.max()}")

    v = np.vstack((lon_vals, lat_vals, values)).T
    verts_df = pd.DataFrame(v, columns=['lon', 'lat', 'var'])
    points = gv.operation.project_points(gv.Points(verts_df, vdims=['var']))
    title = f'{var_name} @ level_idx {level_idx} (surface=0) | Time: {actual_time_selected.strftime("%Y-%m-%d %H:%M")}'
    trimesh = gv.TriMesh((tris_df, points), label=title)
    mesh = dshade.rasterize(trimesh).opts(
        cmap='rainbow',
        colorbar=True,
        width=800,
        height=600,
        clipping_colors={'NaN': 'transparent'}
    )
    print("Plot generated.")
    return tiles * mesh

# ## Create and Display the Dashboard
print("Creating dashboard layout...")
controls = pn.Column(var_widget, level_widget, time_widget, width=300)
dashboard = pn.Row(controls, create_fvcom_plot)

print("Dashboard ready. To view, ensure this script is run in a Jupyter Notebook cell,")
print("or use 'panel serve your_script_name.py' in the terminal and open the browser link.")

# To display in a Jupyter Notebook:
# dashboard

# To make it servable from a script:
# dashboard.servable(title="FVCOM Interactive Explorer")

dashboard