"""Plots for visualizing grid statistics."""

import geopandas as gpd
import matplotlib.colors as mcolors
import numpy as np
import pydeck as pdk

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


def create_grid_areas_map(
    grid_gdf: gpd.GeoDataFrame,
    metric: str,
    make_3d_map: bool,
) -> pdk.Deck:
    """Create a spatial distribution map for grid areas.

    Args:
        grid_gdf (gpd.GeoDataFrame): GeoDataFrame containing grid cell geometries and statistics.
        metric (str): The metric to visualize (e.g., "cell_area", "land_area", "water_area", "land_ratio").
        make_3d_map (bool): Whether to render the map in 3D (extruded) or 2D.

    Returns:
        pdk.Deck: A PyDeck map visualization of the specified grid statistic.

    """
    # Subsample if too many cells for performance
    n_cells = len(grid_gdf)
    if n_cells > 50000:
        grid_gdf = grid_gdf.sample(n=50000, random_state=42)

    # Create a copy to avoid modifying the original
    gdf = grid_gdf.copy().to_crs("EPSG:4326")

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

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

    # Get colormap for the metric
    cmap = get_cmap(metric)

    # Normalize the metric values to [0, 1] for color mapping
    values = gdf_wgs84[metric].values
    vmin, vmax = values.min(), values.max()

    if vmax > vmin:
        normalized_values = (values - vmin) / (vmax - vmin)
    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 metric value for tooltip
    gdf_wgs84["metric_value"] = 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"],
                "metric_value": float(row["metric_value"]),
                "elevation": float(row["elevation"]) if make_3d_map else 0,
                "cell_area": float(row["cell_area"]),
                "land_area": float(row["land_area"]),
                "water_area": float(row["water_area"]),
                "land_ratio": float(row["land_ratio"]),
            },
        }
        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>Cell Area:</b> {cell_area} km²<br/>"
        "<b>Land Area:</b> {land_area} km²<br/>"
        "<b>Water Area:</b> {water_area} km²<br/>"
        "<b>Land Ratio:</b> {land_ratio}"
    )

    # 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
