"""Plots for visualizing ERA5 climate data."""

import geopandas as gpd
import matplotlib.colors as mcolors
import numpy as np
import plotly.graph_objects as go
import pydeck as pdk
import xarray as xr
from plotly.subplots import make_subplots

from entropice.dashboard.utils.colors import get_cmap, hex_to_rgb
from entropice.dashboard.utils.geometry import fix_hex_geometry


def create_climate_map(
    climate_values: xr.DataArray,
    grid_gdf: gpd.GeoDataFrame,
    variable_name: str,
    make_3d_map: bool,
) -> pdk.Deck:
    """Create a spatial distribution map for ERA5 climate variables.

    Args:
        climate_values: Series with cell_ids as index and climate values
        grid_gdf: GeoDataFrame containing grid cell geometries
        variable_name: Name of the climate variable being visualized
        make_3d_map: Whether to render the map in 3D (extruded) or 2D

    Returns:
        pdk.Deck: A PyDeck map visualization of the climate variable

    """
    # Subsample if too many cells for performance
    n_cells = len(climate_values["cell_ids"])
    if n_cells > 50000:
        rng = np.random.default_rng(42)
        cell_indices = rng.choice(n_cells, size=50000, replace=False)
        climate_values = climate_values.isel(cell_ids=cell_indices)
    # Create a copy to avoid modifying the original
    gdf = grid_gdf.copy().to_crs("EPSG:4326")

    # Convert to DataFrame for easier merging
    climate_df = climate_values.to_dataframe(name="climate_value")

    # Reset index if cell_id is already the index
    if gdf.index.name == "cell_id":
        gdf = gdf.reset_index()

    # Filter grid to only cells that have climate data
    gdf = gdf[gdf["cell_id"].isin(climate_df.index)]
    gdf = gdf.set_index("cell_id")

    # Merge climate values with grid geometries
    gdf = gdf.join(climate_df, how="inner")

    # Convert to WGS84 for pydeck
    gdf_wgs84 = gdf.to_crs("EPSG:4326")

    # Fix antimeridian issues for hex cells
    gdf_wgs84["geometry"] = gdf_wgs84["geometry"].apply(fix_hex_geometry)

    # Get colormap
    cmap = get_cmap(variable_name)

    # Normalize the climate values to [0, 1] for color mapping
    values = gdf_wgs84["climate_value"].to_numpy()

    # Use percentiles to avoid outliers
    vmin, vmax = np.nanpercentile(values, [2, 98])
    if vmax > vmin:
        normalized_values = np.clip((values - vmin) / (vmax - vmin), 0, 1)
    else:
        normalized_values = np.zeros_like(values)

    # Map normalized values to colors
    colors = [cmap(val) for val in normalized_values]
    rgb_colors = [hex_to_rgb(mcolors.to_hex(color)) for color in colors]
    gdf_wgs84["fill_color"] = rgb_colors

    # Store climate value for tooltip
    gdf_wgs84["climate_value_display"] = values

    # Store normalized values for elevation (if 3D)
    gdf_wgs84["elevation"] = normalized_values

    # Convert to GeoJSON format
    geojson_data = []
    for _, row in gdf_wgs84.iterrows():
        feature = {
            "type": "Feature",
            "geometry": row["geometry"].__geo_interface__,
            "properties": {
                "fill_color": row["fill_color"],
                "climate_value": float(row["climate_value_display"]),
                "elevation": float(row["elevation"]) if make_3d_map else 0,
            },
        }
        geojson_data.append(feature)

    # Create pydeck layer
    layer = pdk.Layer(
        "GeoJsonLayer",
        geojson_data,
        opacity=0.7,
        stroked=True,
        filled=True,
        extruded=make_3d_map,
        wireframe=False,
        get_fill_color="properties.fill_color",
        get_line_color=[80, 80, 80],
        line_width_min_pixels=0.5,
        get_elevation="properties.elevation" if make_3d_map else 0,
        elevation_scale=500000,
        pickable=True,
    )

    # Set initial view state
    view_state = pdk.ViewState(
        latitude=70,
        longitude=0,
        zoom=2 if not make_3d_map else 1.5,
        pitch=0 if not make_3d_map else 45,
    )

    # Build tooltip HTML
    tooltip_html = f"<b>{variable_name}:</b> {{climate_value}}"

    # Create deck
    deck = pdk.Deck(
        layers=[layer],
        initial_view_state=view_state,
        tooltip={
            "html": tooltip_html,
            "style": {"backgroundColor": "steelblue", "color": "white"},
        },
        map_style="https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json",
    )

    return deck


def create_climate_trend_plot(climate_data: xr.DataArray, variable_name: str) -> go.Figure:
    """Create a trend plot for climate variables over time.

    Args:
        climate_data: DataArray containing climate variable with 'year' dimension
        variable_name: Name of the variable being plotted

    Returns:
        Plotly Figure with trend plot

    """
    # Subsample if too many cells for performance
    n_cells = len(climate_data["cell_ids"])
    if n_cells > 10000:
        rng = np.random.default_rng(42)
        cell_indices = rng.choice(n_cells, size=10000, replace=False)
        climate_data = climate_data.isel({"cell_ids": cell_indices})

    # Get years
    years = climate_data["year"].to_numpy()

    # Calculate statistics over space for each year
    mean_values = climate_data.mean(dim="cell_ids").to_numpy()
    min_values = climate_data.min(dim="cell_ids").to_numpy()
    max_values = climate_data.max(dim="cell_ids").to_numpy()
    p10_values = climate_data.quantile(0.10, dim="cell_ids").to_numpy()
    p90_values = climate_data.quantile(0.90, dim="cell_ids").to_numpy()

    fig = go.Figure()

    # Add min/max range (background)
    fig.add_trace(
        go.Scatter(
            x=years,
            y=min_values,
            mode="lines",
            line={"color": "lightgray", "width": 1, "dash": "dash"},
            name="Min/Max Range",
            showlegend=True,
        )
    )

    fig.add_trace(
        go.Scatter(
            x=years,
            y=max_values,
            mode="lines",
            fill="tonexty",
            fillcolor="rgba(200, 200, 200, 0.1)",
            line={"color": "lightgray", "width": 1, "dash": "dash"},
            showlegend=False,
        )
    )

    # Add 10th-90th percentile band
    fig.add_trace(
        go.Scatter(
            x=years,
            y=p10_values,
            mode="lines",
            line={"width": 0},
            showlegend=False,
            hoverinfo="skip",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=years,
            y=p90_values,
            mode="lines",
            fill="tonexty",
            fillcolor="rgba(33, 150, 243, 0.2)",
            line={"width": 0},
            name="10th-90th Percentile",
        )
    )

    # Add mean line
    fig.add_trace(
        go.Scatter(
            x=years,
            y=mean_values,
            mode="lines+markers",
            name="Mean",
            line={"color": "#1976D2", "width": 2},
            marker={"size": 6},
        )
    )

    fig.update_layout(
        title=f"{variable_name} Over Time (Spatial Statistics)",
        xaxis_title="Year",
        yaxis_title=variable_name,
        height=450,
        hovermode="x unified",
        legend={
            "orientation": "h",
            "yanchor": "bottom",
            "y": 1.02,
            "xanchor": "right",
            "x": 1,
        },
    )

    fig.update_xaxes(dtick=1)

    return fig


def create_climate_distribution_plot(climate_ds: xr.Dataset, variables: list[str]) -> go.Figure:
    """Create distribution plots for climate variables.

    Args:
        climate_ds: Xarray Dataset containing ERA5 climate features
        variables: List of variable names to plot

    Returns:
        Plotly Figure with distribution plots

    """
    # Subsample if too many cells for performance
    n_cells = len(climate_ds.cell_ids)
    if n_cells > 10000:
        rng = np.random.default_rng(42)
        cell_indices = rng.choice(n_cells, size=10000, replace=False)
        climate_ds = climate_ds.isel(cell_ids=cell_indices)

    # Create subplots - one row per variable
    n_rows = len(variables)
    fig = make_subplots(
        rows=n_rows,
        cols=1,
        subplot_titles=[v.replace("_", " ").title() for v in variables],
        vertical_spacing=0.15 / max(n_rows, 1),
    )

    # Color palette
    colors = [
        "#1f77b4",
        "#ff7f0e",
        "#2ca02c",
        "#d62728",
        "#9467bd",
        "#8c564b",
        "#e377c2",
        "#7f7f7f",
    ]

    for row_idx, variable in enumerate(variables, start=1):
        if variable not in climate_ds.data_vars:
            continue

        # Get data
        data = climate_ds[variable]
        values = data.to_numpy().flatten()
        values = values[~np.isnan(values)]

        if len(values) == 0:
            continue

        # Add violin plot
        fig.add_trace(
            go.Violin(
                y=values,
                name=variable.replace("_", " ").title(),
                box_visible=True,
                meanline_visible=True,
                line_color=colors[row_idx % len(colors)],
                showlegend=False,
            ),
            row=row_idx,
            col=1,
        )

    # Update layout
    fig.update_layout(
        height=300 * n_rows,
        title_text="Climate Variable Distributions",
        showlegend=False,
    )

    # Update y-axes labels
    for i in range(n_rows):
        fig.update_yaxes(title_text="Value", row=i + 1, col=1)

    return fig


def create_temperature_comparison_plot(climate_ds: xr.Dataset) -> go.Figure:
    """Create a comparison plot for temperature-related variables.

    Args:
        climate_ds: Xarray Dataset containing temperature variables

    Returns:
        Plotly Figure with temperature comparison

    """
    # Temperature variables to compare
    temp_vars = ["t2m_max", "t2m_mean", "t2m_min"]
    available_vars = [v for v in temp_vars if v in climate_ds.data_vars]

    if not available_vars:
        return go.Figure()

    # Get years
    years = climate_ds["year"].to_numpy()

    fig = go.Figure()

    colors = {"t2m_max": "#d32f2f", "t2m_mean": "#1976d2", "t2m_min": "#0288d1"}

    for var in available_vars:
        # Calculate spatial mean for each year
        values = climate_ds[var].mean(dim="cell_ids").to_numpy()

        fig.add_trace(
            go.Scatter(
                x=years,
                y=values - 273.15,  # Convert to Celsius
                mode="lines+markers",
                name=var.replace("_", " ").title(),
                line={"color": colors.get(var, "#666666"), "width": 2},
                marker={"size": 4},
            )
        )

    # Add freezing point reference line
    fig.add_hline(y=0, line_dash="dash", line_color="gray", annotation_text="Freezing Point (0°C)")

    fig.update_layout(
        title="Temperature Extremes Over Time",
        xaxis_title="Year",
        yaxis_title="Temperature (°C)",
        height=450,
        hovermode="x unified",
    )

    return fig


def create_seasonal_pattern_plot(climate_ds: xr.Dataset, variable: str) -> go.Figure:
    """Create a plot showing seasonal patterns.

    Args:
        climate_ds: Xarray Dataset with month dimension
        variable: Variable name to plot

    Returns:
        Plotly Figure with seasonal patterns

    """
    if variable not in climate_ds.data_vars or "month" not in climate_ds.dims:
        return go.Figure()

    # Get unique months/seasons
    months = climate_ds["month"].to_numpy()

    # Calculate mean across space and years for each month
    values = climate_ds[variable].mean(dim=["cell_ids", "year"]).to_numpy()

    fig = go.Figure()

    fig.add_trace(
        go.Bar(
            x=months,
            y=values,
            marker_color="#1976d2",
        )
    )

    fig.update_layout(
        title=f"{variable.replace('_', ' ').title()} by Season",
        xaxis_title="Season",
        yaxis_title=variable.replace("_", " ").title(),
        height=400,
    )

    return fig
