"""Render the AlphaEarth visualization tab."""

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 entropice.dashboard.utils.colors import get_cmap, hex_to_rgb
from entropice.dashboard.utils.geometry import fix_hex_geometry


def create_embedding_map(
    embedding_values: xr.DataArray,
    grid_gdf: gpd.GeoDataFrame,
    make_3d_map: bool,
) -> pdk.Deck:
    """Create a spatial distribution map for AlphaEarth embeddings.

    Args:
        embedding_values (xr.DataArray): DataArray containing the already filtered AlphaEarth embeddings.
        grid_gdf (gpd.GeoDataFrame): GeoDataFrame containing grid cell geometries.
        make_3d_map (bool): Whether to render the map in 3D (extruded) or 2D.

    Returns:
        pdk.Deck: A PyDeck map visualization of the AlphaEarth embeddings.

    """
    # Subsample if too many cells for performance
    n_cells = len(embedding_values["cell_ids"])
    if n_cells > 50000:
        rng = np.random.default_rng(42)  # Fixed seed for reproducibility
        cell_indices = rng.choice(n_cells, size=50000, replace=False)
        embedding_values = embedding_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
    embedding_df = embedding_values.to_dataframe(name="embedding_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 embedding data
    gdf = gdf[gdf["cell_id"].isin(embedding_df.index)]
    gdf = gdf.set_index("cell_id")

    # Merge embedding values with grid geometries
    gdf = gdf.join(embedding_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 for embeddings
    cmap = get_cmap("AlphaEarth")

    # Normalize the embedding values to [0, 1] for color mapping
    # Use percentiles to avoid outliers
    values = gdf_wgs84["embedding_value"].values
    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 embedding value for tooltip
    gdf_wgs84["embedding_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"],
                "embedding_value": float(row["embedding_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,  # Scale normalized values (0-1) to 500km height
        pickable=True,
    )

    # Set initial view state (centered on the Arctic)
    # Adjust pitch and zoom based on whether we're using 3D
    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 = "<b>Embedding Value:</b> {embedding_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_embedding_trend_plot(embedding_values: xr.DataArray) -> go.Figure:
    """Create a trend plot for AlphaEarth embeddings over time.

    Contains a line plot with shaded areas representing the 10th to 90th percentiles.
    Min and Max values are marked through a dashed line.

    Args:
        embedding_values (xr.DataArray): DataArray containing the AlphaEarth embeddings with a 'year' dimension.

    Returns:
        go.Figure: A Plotly figure showing the trend of embeddings over time.

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

    # Calculate statistics over space (cell_ids) for each year
    years = embedding_values["year"].values

    # Compute statistics across cells for each year
    mean_values = embedding_values.mean(dim="cell_ids").to_numpy()
    min_values = embedding_values.min(dim="cell_ids").to_numpy()
    max_values = embedding_values.max(dim="cell_ids").to_numpy()
    p10_values = embedding_values.quantile(0.10, dim="cell_ids").to_numpy()
    p90_values = embedding_values.quantile(0.90, dim="cell_ids").to_numpy()

    fig = go.Figure()

    # Add min/max range first (background) - dashed lines
    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(76, 175, 80, 0.2)",  # Green shade for AlphaEarth
            line={"width": 0},
            name="10th-90th Percentile",
        )
    )

    # Add mean line on top
    fig.add_trace(
        go.Scatter(
            x=years,
            y=mean_values,
            mode="lines+markers",
            name="Mean",
            line={"color": "#2E7D32", "width": 2},  # Darker green for AlphaEarth
            marker={"size": 6},
        )
    )

    fig.update_layout(
        title="Embedding Values Over Time (Spatial Statistics)",
        xaxis_title="Year",
        yaxis_title="Embedding Value",
        yaxis={"zeroline": True, "zerolinewidth": 2, "zerolinecolor": "gray"},
        height=450,
        hovermode="x unified",
        legend={
            "orientation": "h",
            "yanchor": "bottom",
            "y": 1.02,
            "xanchor": "right",
            "x": 1,
        },
    )

    # Format x-axis to show years as integers
    fig.update_xaxes(dtick=1)

    return fig


def create_embedding_distribution_plot(embedding_values: xr.DataArray) -> go.Figure:
    """Create a distribution plot showing the min/max, 10th/90th percentiles, and mean of AlphaEarth bands.

    Args:
        embedding_values (xr.DataArray): DataArray containing the AlphaEarth embeddings.

    Returns:
        go.Figure: A Plotly figure showing the distribution of embeddings.

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

    # Get band dimension
    bands = embedding_values["band"].values

    # Calculate statistics for each band across all cells
    band_stats = []
    for band in bands:
        band_data = embedding_values.sel(band=band).values.flatten()
        # Remove NaN values
        band_data = band_data[~np.isnan(band_data)]

        if len(band_data) > 0:
            band_stats.append(
                {
                    "Band": str(band),
                    "Mean": float(np.mean(band_data)),
                    "Min": float(np.min(band_data)),
                    "Max": float(np.max(band_data)),
                    "P10": float(np.percentile(band_data, 10)),
                    "P90": float(np.percentile(band_data, 90)),
                }
            )

    # Create DataFrame from statistics
    import pandas as pd

    band_df = pd.DataFrame(band_stats)

    fig = go.Figure()

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

    fig.add_trace(
        go.Scatter(
            x=band_df["Band"],
            y=band_df["Max"],
            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=band_df["Band"],
            y=band_df["P10"],
            mode="lines",
            line={"width": 0},
            showlegend=False,
            hoverinfo="skip",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=band_df["Band"],
            y=band_df["P90"],
            mode="lines",
            fill="tonexty",
            fillcolor="rgba(76, 175, 80, 0.2)",  # Green shade for AlphaEarth
            line={"width": 0},
            name="10th-90th Percentile",
        )
    )

    # Add mean line on top
    fig.add_trace(
        go.Scatter(
            x=band_df["Band"],
            y=band_df["Mean"],
            mode="lines+markers",
            name="Mean",
            line={"color": "#2E7D32", "width": 2},  # Darker green for AlphaEarth
            marker={"size": 4},
        )
    )

    fig.update_layout(
        title="Embedding Distribution by Band (Spatial Statistics)",
        xaxis_title="Band",
        yaxis_title="Embedding Value",
        height=450,
        hovermode="x unified",
        legend={
            "orientation": "h",
            "yanchor": "bottom",
            "y": 1.02,
            "xanchor": "right",
            "x": 1,
        },
    )

    return fig
