"""Plots for visualizing ArcticDEM terrain features."""

import geopandas as gpd
import matplotlib.colors as mcolors
import numpy as np
import pandas as pd
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_terrain_map(
    terrain_values: pd.Series,
    grid_gdf: gpd.GeoDataFrame,
    variable_name: str,
    make_3d_map: bool,
) -> pdk.Deck:
    """Create a spatial distribution map for ArcticDEM terrain features.

    Args:
        terrain_values: Series with cell_ids as index and terrain values
        grid_gdf: GeoDataFrame containing grid cell geometries
        variable_name: Name of the terrain 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 terrain feature

    """
    # Subsample if too many cells for performance
    n_cells = len(terrain_values)
    if n_cells > 50000:
        rng = np.random.default_rng(42)
        cell_indices = rng.choice(n_cells, size=50000, replace=False)
        terrain_values = terrain_values.iloc[cell_indices]

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

    # 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 terrain data
    gdf = gdf[gdf["cell_id"].isin(terrain_values.index)]
    gdf = gdf.set_index("cell_id")

    # Merge terrain values with grid geometries
    gdf = gdf.join(terrain_values.to_frame("terrain_value"), 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 - use special colormap for aspect (circular)
    if "aspect" in variable_name.lower():
        cmap = get_cmap("aspect")
    else:
        cmap = get_cmap("terrain")

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

    # Handle aspect specially (0-360 degrees, circular)
    if "aspect" in variable_name.lower():
        vmin, vmax = 0, 360
        normalized_values = values / 360
    else:
        # 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 terrain value for tooltip
    gdf_wgs84["terrain_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"],
                "terrain_value": float(row["terrain_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> {{terrain_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_terrain_distribution_plot(arcticdem_ds: xr.Dataset, features: list[str]) -> go.Figure:
    """Create distribution plots for terrain features.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features
        features: List of feature names to plot

    Returns:
        Plotly Figure with distribution plots

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

    # Determine aggregation types available
    aggs = list(arcticdem_ds.coords["aggregations"].values)

    # Create subplots - one row per aggregation type
    n_rows = len(aggs)
    fig = make_subplots(
        rows=n_rows,
        cols=1,
        subplot_titles=[f"{agg.title()} Values" for agg in aggs],
        vertical_spacing=0.15 / max(n_rows, 1),
    )

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

    for row_idx, agg in enumerate(aggs, start=1):
        for feat_idx, feature in enumerate(features):
            # Get the data for this feature and aggregation
            var_name = f"{feature}"
            if var_name not in arcticdem_ds.data_vars:
                continue

            # Extract values for this aggregation
            data = arcticdem_ds[var_name].sel(aggregations=agg)
            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=feature.replace("_", " ").title(),
                    box_visible=True,
                    meanline_visible=True,
                    line_color=colors[feat_idx % len(colors)],
                    showlegend=(row_idx == 1),
                ),
                row=row_idx,
                col=1,
            )

    # Update layout
    fig.update_layout(
        height=300 * n_rows,
        title_text="Terrain Feature Distributions by Aggregation Type",
        showlegend=True,
    )

    # 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_aspect_rose_diagram(aspect_values: np.ndarray) -> go.Figure:
    """Create a rose diagram (circular histogram) for aspect values.

    Args:
        aspect_values: Array of aspect values in degrees (0-360)

    Returns:
        Plotly Figure with rose diagram

    """
    # Remove NaN values
    aspect_values = aspect_values[~np.isnan(aspect_values)]

    # Subsample if too many values for performance
    if len(aspect_values) > 50000:
        rng = np.random.default_rng(42)
        indices = rng.choice(len(aspect_values), size=50000, replace=False)
        aspect_values = aspect_values[indices]

    if len(aspect_values) == 0:
        # Return empty figure
        return go.Figure()

    # Create bins for aspect (every 10 degrees)
    bins = np.arange(0, 361, 10)
    bin_counts, _ = np.histogram(aspect_values, bins=bins)

    # Calculate bin centers in degrees and radians
    bin_centers_deg = (bins[:-1] + bins[1:]) / 2
    bin_centers_rad = np.deg2rad(bin_centers_deg)

    # Close the circle
    bin_centers_rad = np.append(bin_centers_rad, bin_centers_rad[0])
    bin_counts = np.append(bin_counts, bin_counts[0])

    # Create polar bar chart
    fig = go.Figure()

    fig.add_trace(
        go.Barpolar(
            r=bin_counts,
            theta=np.rad2deg(bin_centers_rad),
            width=10,
            marker={
                "color": bin_centers_rad,
                "colorscale": "HSV",
                "cmin": 0,
                "cmax": 2 * np.pi,
                "showscale": False,
            },
        )
    )

    fig.update_layout(
        title="Aspect Rose Diagram",
        polar={
            "radialaxis": {"title": "Frequency", "showticklabels": True},
            "angularaxis": {
                "direction": "clockwise",
                "rotation": 90,
                "tickmode": "array",
                "tickvals": [0, 45, 90, 135, 180, 225, 270, 315],
                "ticktext": ["N", "NE", "E", "SE", "S", "SW", "W", "NW"],
            },
        },
        showlegend=False,
        height=500,
    )

    return fig


def create_slope_aspect_scatter(slope_values: np.ndarray, aspect_values: np.ndarray) -> go.Figure:
    """Create a scatter plot showing the relationship between slope and aspect.

    Args:
        slope_values: Array of slope values
        aspect_values: Array of aspect values in degrees (0-360)

    Returns:
        Plotly Figure with scatter plot

    """
    # Create DataFrame and remove NaN
    df = pd.DataFrame({"slope": slope_values.flatten(), "aspect": aspect_values.flatten()})
    df = df.dropna()

    if len(df) == 0:
        return go.Figure()

    # Sample if too many points for performance
    if len(df) > 50000:
        df = df.sample(n=50000, random_state=42)

    # Create 2D histogram (density plot)
    fig = go.Figure()

    fig.add_trace(
        go.Histogram2d(
            x=df["aspect"],
            y=df["slope"],
            colorscale="Viridis",
            nbinsx=36,  # 10-degree bins for aspect
            nbinsy=50,
        )
    )

    fig.update_layout(
        title="Slope vs Aspect Distribution",
        xaxis_title="Aspect (degrees)",
        yaxis_title="Slope (degrees)",
        height=500,
    )

    # Add directional labels on x-axis
    fig.update_xaxes(
        tickmode="array",
        tickvals=[0, 45, 90, 135, 180, 225, 270, 315, 360],
        ticktext=["N", "NE", "E", "SE", "S", "SW", "W", "NW", "N"],
    )

    return fig


def create_correlation_heatmap(arcticdem_ds: xr.Dataset, features: list[str], agg: str) -> go.Figure:
    """Create a correlation heatmap for terrain features.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features
        features: List of feature names to include
        agg: Aggregation type to use

    Returns:
        Plotly Figure with correlation heatmap

    """
    # Extract data for each feature
    data_dict = {}
    for feature in features:
        var_name = f"{feature}"
        if var_name in arcticdem_ds.data_vars:
            data = arcticdem_ds[var_name].sel(aggregations=agg)
            values = data.to_numpy().flatten()
            data_dict[feature] = values

    if not data_dict:
        return go.Figure()

    # Create DataFrame
    df = pd.DataFrame(data_dict)
    df = df.dropna()

    # Sample if too many rows
    if len(df) > 50000:
        df = df.sample(n=50000, random_state=42)

    # Calculate correlation matrix
    corr = df.corr()

    # Create heatmap
    fig = go.Figure(
        data=go.Heatmap(
            z=corr.values,
            x=[f.replace("_", " ").title() for f in corr.columns],
            y=[f.replace("_", " ").title() for f in corr.index],
            colorscale="RdBu",
            zmid=0,
            zmin=-1,
            zmax=1,
            text=np.round(corr.values, 2),
            texttemplate="%{text}",
            textfont={"size": 10},
            colorbar={"title": "Correlation"},
        )
    )

    fig.update_layout(
        title=f"Feature Correlation Matrix ({agg.title()})",
        height=600,
        xaxis={"side": "bottom"},
    )

    return fig
