"""Plots for visualizing model inference results."""

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

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


def create_inference_map(
    predictions_gdf: gpd.GeoDataFrame,
    task: Task,
    make_3d_map: bool,
) -> pdk.Deck:
    """Create a spatial distribution map for model predictions.

    Args:
        predictions_gdf: GeoDataFrame with columns ['cell_id', 'predicted', 'geometry']
        task: Task type (binary, count_regimes, density_regimes, count, density)
        make_3d_map: Whether to render the map in 3D (extruded) or 2D

    Returns:
        pdk.Deck: A PyDeck map visualization of the predictions

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

    # Create a copy to avoid modifying the original
    gdf = predictions_gdf.copy()

    # Convert to WGS84 for pydeck
    if gdf.crs is not None:
        gdf = gdf.to_crs("EPSG:4326")

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

    # Determine if classification or regression
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]
    is_classification = task in classification_tasks

    if is_classification:
        # Convert to categorical if needed
        if gdf["predicted"].dtype != "category":
            gdf["predicted"] = pd.Categorical(gdf["predicted"])

        # Get categories and assign colors
        categories = gdf["predicted"].cat.categories.tolist()
        n_categories = len(categories)
        colors = get_palette(task, n_categories + 2)[1:-1]  # Avoid too light/dark colors

        # Create category to color mapping
        category_colors = {cat: colors[i] for i, cat in enumerate(categories)}

        # Assign colors to each row based on predicted category
        rgb_colors = []
        for cat in gdf["predicted"]:
            color = category_colors[cat]
            rgb_colors.append(hex_to_rgb(mcolors.to_hex(color)))
        gdf["fill_color"] = rgb_colors

        # For elevation in 3D, use category index normalized
        if n_categories > 1:
            gdf["elevation"] = gdf["predicted"].cat.codes / (n_categories - 1)
        else:
            gdf["elevation"] = 0.5

        # Display value is the category name
        gdf["display_value"] = gdf["predicted"].astype(str)

    else:
        # Regression task - use continuous colormap
        cmap = get_cmap(task)

        # Get prediction values
        values = gdf["predicted"].to_numpy()

        # Normalize using percentiles to avoid outliers
        vmin, vmax = np.nanpercentile(values, [2, 98])
        if vmax > vmin:
            normalized_values = (np.clip(values, vmin, vmax) - vmin) / (vmax - vmin)
        else:
            normalized_values = np.full_like(values, 0.5)

        # 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["fill_color"] = rgb_colors

        # Use normalized values for elevation
        gdf["elevation"] = normalized_values

        # Display value is the actual prediction (rounded)
        gdf["display_value"] = gdf["predicted"].round(3).astype(str)

    # Convert to GeoJSON format
    geojson_data = []
    for _, row in gdf.iterrows():
        geojson_data.append(
            {
                "type": "Feature",
                "geometry": row["geometry"].__geo_interface__,
                "properties": {
                    "fill_color": row["fill_color"],
                    "elevation": row["elevation"],
                    "predicted": row["display_value"],
                },
            }
        )

    # 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 = "<b>Predicted:</b> {predicted}"

    # 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_prediction_distribution_plot(
    predictions_gdf: gpd.GeoDataFrame,
    task: Task,
) -> go.Figure:
    """Create a distribution plot for predictions.

    For classification tasks, shows a bar chart of class counts.
    For regression tasks, shows a histogram of predicted values.

    Args:
        predictions_gdf: GeoDataFrame with 'predicted' column
        task: Task type

    Returns:
        Plotly Figure with distribution plot

    """
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]
    is_classification = task in classification_tasks

    fig = go.Figure()

    if is_classification:
        # Bar chart for classification
        # Convert to categorical if needed
        if predictions_gdf["predicted"].dtype != "category":
            predictions_gdf["predicted"] = pd.Categorical(predictions_gdf["predicted"])

        categories = predictions_gdf["predicted"].cat.categories.tolist()
        counts = predictions_gdf["predicted"].value_counts().reindex(categories, fill_value=0)

        # Get colors
        n_categories = len(categories)
        colors = get_palette(task, n_categories + 2)[1:-1]

        fig.add_trace(
            go.Bar(
                x=categories,
                y=counts.to_numpy(),
                marker_color=colors,
                text=counts.to_numpy(),
                textposition="outside",
                textfont={"size": 12},
                hovertemplate="<b>%{x}</b><br>Count: %{y}<br>Percentage: %{customdata:.1%}<extra></extra>",
                customdata=counts.to_numpy() / counts.sum(),
            )
        )

        fig.update_layout(
            title=f"Predicted Class Distribution ({task.replace('_', ' ').title()})",
            xaxis_title="Predicted Class",
            yaxis_title="Count",
            height=500,
            showlegend=False,
        )

        fig.update_xaxes(tickangle=-45)

    else:
        # Histogram for regression
        values = predictions_gdf["predicted"].to_numpy()

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

        # Create histogram
        fig.add_trace(
            go.Histogram(
                x=values,
                nbinsx=50,
                marker_color=mcolors.to_hex(get_cmap(task)(0.6)),
                opacity=0.7,
                hovertemplate="<b>Range:</b> %{x}<br><b>Count:</b> %{y}<extra></extra>",
            )
        )

        # Add statistical lines
        mean_val = np.mean(values)
        median_val = np.median(values)

        fig.add_vline(
            x=mean_val,
            line_dash="dash",
            line_color="red",
            annotation_text=f"Mean: {mean_val:.3f}",
            annotation_position="top",
        )

        fig.add_vline(
            x=median_val,
            line_dash="dot",
            line_color="green",
            annotation_text=f"Median: {median_val:.3f}",
            annotation_position="bottom right",
        )

        fig.update_layout(
            title=f"Predicted Value Distribution ({task.replace('_', ' ').title()})",
            xaxis_title=f"Predicted {task.capitalize()}",
            yaxis_title="Count",
            height=500,
            showlegend=False,
        )

    return fig


def create_prediction_statistics_plot(
    predictions_gdf: gpd.GeoDataFrame,
    task: Task,
) -> go.Figure:
    """Create a plot showing key statistics about the predictions.

    Args:
        predictions_gdf: GeoDataFrame with 'predicted' column
        task: Task type

    Returns:
        Plotly Figure with statistics

    """
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]
    is_classification = task in classification_tasks

    if is_classification:
        # Classification statistics: class proportions
        # Convert to categorical if needed
        if predictions_gdf["predicted"].dtype != "category":
            predictions_gdf["predicted"] = pd.Categorical(predictions_gdf["predicted"])

        categories = predictions_gdf["predicted"].cat.categories.tolist()
        counts = predictions_gdf["predicted"].value_counts().reindex(categories, fill_value=0)
        proportions = counts / counts.sum()

        # Get colors
        n_categories = len(categories)
        colors = get_palette(task, n_categories + 2)[1:-1]

        fig = go.Figure()

        fig.add_trace(
            go.Bar(
                x=categories,
                y=proportions.to_numpy() * 100,  # Convert to percentage
                marker_color=colors,
                text=[f"{p:.1%}" for p in proportions.to_numpy()],
                textposition="outside",
                hovertemplate="<b>%{x}</b><br>Proportion: %{y:.1f}%<extra></extra>",
            )
        )

        fig.update_layout(
            title="Predicted Class Proportions",
            xaxis_title="Predicted Class",
            yaxis_title="Proportion (%)",
            height=400,
            showlegend=False,
        )

        fig.update_xaxes(tickangle=-45)

    else:
        # Regression statistics: box plot with quartiles
        values = predictions_gdf["predicted"].to_numpy()
        values = values[~np.isnan(values)]

        fig = go.Figure()

        fig.add_trace(
            go.Box(
                y=values,
                name="Predictions",
                marker_color=mcolors.to_hex(get_cmap(task)(0.6)),
                boxmean="sd",  # Show mean and standard deviation
                hovertemplate="<b>Value:</b> %{y:.3f}<extra></extra>",
            )
        )

        # Calculate statistics
        stats = pd.Series(values).describe()

        # Add annotation with statistics
        stats_text = (
            f"<b>Statistics</b><br>"
            f"Count: {int(stats['count'])}<br>"
            f"Mean: {stats['mean']:.3f}<br>"
            f"Std: {stats['std']:.3f}<br>"
            f"Min: {stats['min']:.3f}<br>"
            f"25%: {stats['25%']:.3f}<br>"
            f"50%: {stats['50%']:.3f}<br>"
            f"75%: {stats['75%']:.3f}<br>"
            f"Max: {stats['max']:.3f}"
        )

        fig.add_annotation(
            text=stats_text,
            xref="paper",
            yref="paper",
            x=1.15,
            y=0.5,
            showarrow=False,
            align="left",
            bgcolor="rgba(255, 255, 255, 0.8)",
            bordercolor="black",
            borderwidth=1,
        )

        fig.update_layout(
            title=f"Prediction Statistics ({task.replace('_', ' ').title()})",
            yaxis_title=f"Predicted {task.capitalize()}",
            height=500,
            showlegend=False,
            margin={"r": 200},  # Make room for stats annotation
        )

    return fig


def create_coverage_statistics_table(
    predictions_gdf: gpd.GeoDataFrame, task: Task, grid_gdf: gpd.GeoDataFrame | None = None
) -> pd.DataFrame:
    """Create a summary table of prediction coverage statistics.

    Args:
        predictions_gdf: GeoDataFrame with predictions
        task: Task type
        grid_gdf: Optional GeoDataFrame with cell_id and land_area columns for RTS area calculations

    Returns:
        DataFrame with coverage statistics

    """
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]
    is_classification = task in classification_tasks

    stats = {
        "Total Cells": len(predictions_gdf),
        "Non-null Predictions": predictions_gdf["predicted"].notna().sum(),
        "Null Predictions": predictions_gdf["predicted"].isna().sum(),
    }

    # Calculate RTS-affected area if grid_gdf is provided
    if grid_gdf is not None:
        # Merge with grid to get land areas
        merged = predictions_gdf.merge(grid_gdf[["cell_id", "land_area"]], on="cell_id", how="left")

        if is_classification:
            # Convert to categorical if needed
            if merged["predicted"].dtype != "category":
                merged["predicted"] = pd.Categorical(merged["predicted"])

            categories = merged["predicted"].cat.categories.tolist()

            # For classification: sum areas where RTS are expected (non-zero classes)
            # Assuming first category is "no RTS" (0 or similar)
            rts_mask = merged["predicted"].notna()
            if len(categories) > 1:
                # Exclude the first category (assumed to be "no RTS")
                rts_mask = rts_mask & (merged["predicted"] != categories[0])

            rts_area_km2 = merged.loc[rts_mask, "land_area"].sum() / 1e6  # Convert m² to km²
            stats["RTS-Affected Area (km²)"] = f"{rts_area_km2:,.2f}"

        elif task == "count":
            # For count: sum areas where count > 0
            rts_mask = (merged["predicted"].notna()) & (merged["predicted"] > 0)
            rts_area_km2 = merged.loc[rts_mask, "land_area"].sum() / 1e6  # Convert m² to km²
            stats["RTS-Affected Area (km²)"] = f"{rts_area_km2:,.2f}"

        elif task == "density":
            # For density: multiply predicted density by area
            # Density is typically RTS/km², so multiply by area in km²
            merged["rts_area"] = merged["predicted"] * merged["land_area"] / 1e6  # density * area_km2
            total_rts_area = merged["rts_area"].sum()
            stats["Total RTS Area (km²)"] = f"{total_rts_area:,.2f}"

    if is_classification:
        # Add class counts
        # Convert to categorical if needed
        if predictions_gdf["predicted"].dtype != "category":
            predictions_gdf["predicted"] = pd.Categorical(predictions_gdf["predicted"])

        categories = predictions_gdf["predicted"].cat.categories.tolist()
        for cat in categories:
            count = (predictions_gdf["predicted"] == cat).sum()
            stats[f"Class '{cat}' Count"] = count
            stats[f"Class '{cat}' Proportion"] = f"{count / len(predictions_gdf):.2%}"
    else:
        # Add regression statistics
        values = predictions_gdf["predicted"].dropna()
        if len(values) > 0:
            stats["Mean"] = f"{values.mean():.4f}"
            stats["Median"] = f"{values.median():.4f}"
            stats["Std Dev"] = f"{values.std():.4f}"
            stats["Min"] = f"{values.min():.4f}"
            stats["Max"] = f"{values.max():.4f}"

    # Convert to DataFrame
    df = pd.DataFrame(list(stats.items()), columns=["Statistic", "Value"])
    return df
