"""Plots for experiment comparison and analysis."""

from typing import TYPE_CHECKING

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

from entropice.dashboard.utils.colors import get_palette
from entropice.dashboard.utils.formatters import format_metric_name

if TYPE_CHECKING:
    import geopandas as gpd
    import pydeck as pdk


def create_grid_level_comparison_plot(
    results_df: pd.DataFrame,
    metric: str,
    split: str = "test",
) -> go.Figure:
    """Create a plot comparing model performance across grid levels.

    Args:
        results_df: DataFrame with experiment results including grid, level, model, and metrics
        metric: Metric to compare (e.g., 'f1', 'accuracy', 'r2')
        split: Data split to show ('train', 'test', or 'complete')

    Returns:
        Plotly figure showing performance across grid levels

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    # Create grid_level column for grouping
    results_df = results_df.copy()
    results_df["grid_level"] = results_df["grid"] + "_" + results_df["level"].astype(str)

    # Define the proper order for grid levels by resolution
    grid_level_order = [
        "hex_3",
        "healpix_6",
        "healpix_7",
        "hex_4",
        "healpix_8",
        "hex_5",
        "healpix_9",
        "healpix_10",
        "hex_6",
    ]

    # Create display labels for grid levels
    grid_level_display = {
        "hex_3": "Hex-3",
        "healpix_6": "Healpix-6",
        "healpix_7": "Healpix-7",
        "hex_4": "Hex-4",
        "healpix_8": "Healpix-8",
        "hex_5": "Hex-5",
        "healpix_9": "Healpix-9",
        "healpix_10": "Healpix-10",
        "hex_6": "Hex-6",
    }

    # Add display column
    results_df["grid_level_display"] = results_df["grid_level"].map(grid_level_display)

    # Create color map for target datasets (use 2nd color from 5-color palette for more saturation)
    unique_targets = results_df["target"].unique()
    target_colors = {target: get_palette(target, 5)[1] for target in unique_targets}

    # Create symbol map for models
    model_symbols = {
        "espa": "circle",
        "xgboost": "square",
        "rf": "diamond",
        "knn": "cross",
        "ensemble": "star",
        "autogluon": "star",
    }

    # Add a complete column for hover information
    results_df["model_display"] = results_df["model"].str.upper()

    # Create box plot without individual points first
    fig = px.box(
        results_df,
        x="grid_level_display",
        y=metric_col,
        color="target",
        facet_col="task",
        points=False,  # We'll add points separately with symbols
        title=f"{format_metric_name(metric)} by Grid Level ({split.capitalize()} Set)",
        labels={
            metric_col: format_metric_name(metric),
            "grid_level_display": "Grid Level",
            "target": "Target Dataset",
        },
        color_discrete_map=target_colors,
        category_orders={"grid_level_display": [grid_level_display[gl] for gl in grid_level_order]},
    )

    # Add scatter points with model-specific symbols
    # Group by task for faceting
    if "task" in results_df.columns:
        unique_tasks = results_df["task"].unique()
    else:
        unique_tasks = [None]

    # Add scatter traces and line traces for each target-model combination
    for target in unique_targets:
        target_data = results_df[results_df["target"] == target]
        for model in target_data["model"].unique():
            model_data = target_data[target_data["model"] == model]
            symbol = model_symbols.get(model, "circle")

            # Add scatter trace for each task facet
            for task_idx, task in enumerate(unique_tasks):
                if task is not None:
                    task_data = model_data[model_data["task"] == task]
                else:
                    task_data = model_data

                if len(task_data) == 0:
                    continue

                # Sort by grid_level order for proper line connections
                task_data["grid_level_order"] = task_data["grid_level"].map(
                    {gl: i for i, gl in enumerate(grid_level_order)}
                )
                task_data = task_data.sort_values("grid_level_order")

                # Determine which subplot this goes to
                row = 1
                col = task_idx + 1 if task is not None else 1

                # Add line trace connecting points of the same model
                line = go.Scatter(
                    x=task_data["grid_level_display"],
                    y=task_data[metric_col],
                    mode="lines",
                    line={
                        "color": target_colors[target],
                        "width": 1.5,
                        "dash": "dot",
                    },
                    showlegend=False,
                    hoverinfo="skip",
                    xaxis=f"x{col}" if col > 1 else "x",
                    yaxis=f"y{col}" if col > 1 else "y",
                )

                fig.add_trace(line, row=row, col=col)

                # Add scatter trace on top of the line
                scatter = go.Scatter(
                    x=task_data["grid_level_display"],
                    y=task_data[metric_col],
                    mode="markers",
                    marker={
                        "symbol": symbol,
                        "size": 8,
                        "color": target_colors[target],
                        "line": {"width": 1, "color": "white"},
                    },
                    name=f"{target} ({model.upper()})",
                    legendgroup=target,
                    showlegend=(task_idx == 0),  # Only show in legend once
                    hovertemplate=(
                        f"<b>{target}</b><br>"
                        f"Model: {model.upper()}<br>"
                        f"Grid Level: %{{x}}<br>"
                        f"{format_metric_name(metric)}: %{{y:.4f}}<br>"
                        "<extra></extra>"
                    ),
                    xaxis=f"x{col}" if col > 1 else "x",
                    yaxis=f"y{col}" if col > 1 else "y",
                )

                fig.add_trace(scatter, row=row, col=col)

    fig.update_layout(
        height=500,
        showlegend=True,
        hovermode="closest",
    )

    return fig


def create_model_ranking_plot(
    results_df: pd.DataFrame,
    metric: str,
    split: str = "test",
    top_n: int = 10,
) -> go.Figure:
    """Create a plot ranking models by performance.

    Args:
        results_df: DataFrame with experiment results
        metric: Metric to rank by
        split: Data split to show
        top_n: Number of top models to show

    Returns:
        Plotly figure showing top models

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    # Get top N models
    top_models = results_df.nlargest(top_n, metric_col)

    # Create display label
    top_models = top_models.copy()
    top_models["model_label"] = (
        top_models["model"].astype(str)
        + " ("
        + top_models["grid"]
        + "_"
        + top_models["level"].astype(str)
        + ", "
        + top_models["task"]
        + ")"
    )

    colors = get_palette("models", len(top_models))

    fig = go.Figure(
        data=[
            go.Bar(
                x=top_models[metric_col],
                y=top_models["model_label"],
                orientation="h",
                marker={"color": colors},
                text=top_models[metric_col].round(4),
                textposition="auto",
            )
        ]
    )

    fig.update_layout(
        title=f"Top {top_n} Models by {format_metric_name(metric)} ({split.capitalize()} Set)",
        xaxis_title=format_metric_name(metric),
        yaxis_title="Model Configuration",
        height=max(400, top_n * 40),
        showlegend=False,
    )

    return fig


def create_model_consistency_heatmap(
    results_df: pd.DataFrame,
    metric: str,
    split: str = "test",
) -> go.Figure:
    """Create a heatmap showing model consistency across tasks and targets.

    Args:
        results_df: DataFrame with experiment results
        metric: Metric to analyze
        split: Data split to show

    Returns:
        Plotly figure showing heatmap of model performance

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    # Create a pivot table: models vs (task, target)
    results_df = results_df.copy()
    results_df["task_target"] = results_df["task"] + "_" + results_df["target"]

    # Get best score per model-task-target combination
    pivot_data = results_df.groupby(["model", "task_target"])[metric_col].max().reset_index()
    pivot_table = pivot_data.pivot_table(index="model", columns="task_target", values=metric_col)

    fig = go.Figure(
        data=go.Heatmap(
            z=pivot_table.to_numpy(),
            x=pivot_table.columns,
            y=pivot_table.index,
            colorscale="Viridis",
            text=pivot_table.to_numpy().round(3),
            texttemplate="%{text}",
            textfont={"size": 10},
            colorbar={"title": format_metric_name(metric)},
        )
    )

    fig.update_layout(
        title=f"Model Consistency: {format_metric_name(metric)} Across Tasks and Targets ({split.capitalize()} Set)",
        xaxis_title="Task_Target",
        yaxis_title="Model",
        height=max(400, len(pivot_table.index) * 50),
    )

    return fig


def create_feature_importance_comparison_plot(
    feature_importance_df: pd.DataFrame,
    top_n: int = 20,
) -> go.Figure:
    """Create a plot comparing feature importance across different models/configurations.

    Args:
        feature_importance_df: DataFrame with columns: feature, importance, model, grid, level, task
        top_n: Number of top features to show

    Returns:
        Plotly figure showing feature importance comparison

    """
    # Get top features overall
    overall_importance = feature_importance_df.groupby("feature")["importance"].mean().reset_index()
    top_features = overall_importance.nlargest(top_n, "importance")["feature"].tolist()

    # Filter to top features
    filtered_df = feature_importance_df[feature_importance_df["feature"].isin(top_features)]

    # Create grid_level column
    filtered_df = filtered_df.copy()
    filtered_df["config"] = (
        filtered_df["model"].astype(str)
        + " ("
        + filtered_df["grid"]
        + "_"
        + filtered_df["level"].astype(str)
        + ", "
        + filtered_df["task"]
        + ")"
    )

    # Create grouped bar chart
    fig = px.bar(
        filtered_df,
        x="feature",
        y="importance",
        color="config",
        barmode="group",
        title=f"Top {top_n} Features: Importance Comparison Across Configurations",
        labels={"importance": "Feature Importance", "feature": "Feature", "config": "Configuration"},
    )

    fig.update_layout(
        height=600,
        xaxis_tickangle=-45,
        showlegend=True,
        legend={"orientation": "v", "yanchor": "top", "y": 1, "xanchor": "left", "x": 1.02},
    )

    return fig


def create_feature_importance_heatmap(
    feature_importance_df: pd.DataFrame,
    top_n: int = 30,
) -> go.Figure:
    """Create a heatmap of feature importance across configurations.

    Args:
        feature_importance_df: DataFrame with columns: feature, importance, model, grid, level, task
        top_n: Number of top features to show

    Returns:
        Plotly figure showing heatmap

    """
    # Get top features overall
    overall_importance = feature_importance_df.groupby("feature")["importance"].mean().reset_index()
    top_features = overall_importance.nlargest(top_n, "importance")["feature"].tolist()

    # Filter to top features
    filtered_df = feature_importance_df[feature_importance_df["feature"].isin(top_features)]

    # Create config column
    filtered_df = filtered_df.copy()
    filtered_df["config"] = (
        filtered_df["grid"]
        + "_"
        + filtered_df["level"].astype(str)
        + "_"
        + filtered_df["task"]
        + "_"
        + filtered_df["model"].astype(str)
    )

    # Pivot: features vs configs
    pivot_table = filtered_df.pivot_table(
        index="feature", columns="config", values="importance", aggfunc="mean", fill_value=0
    )

    fig = go.Figure(
        data=go.Heatmap(
            z=pivot_table.to_numpy(),
            x=pivot_table.columns,
            y=pivot_table.index,
            colorscale="Viridis",
            colorbar={"title": "Importance"},
        )
    )

    fig.update_layout(
        title=f"Top {top_n} Features: Importance Heatmap Across Configurations",
        xaxis_title="Configuration",
        yaxis_title="Feature",
        height=max(500, len(pivot_table.index) * 20),
        xaxis_tickangle=-45,
    )

    return fig


def create_model_performance_distribution(
    results_df: pd.DataFrame,
    metric: str,
    split: str = "test",
) -> go.Figure:
    """Create violin plots showing performance distribution by model type.

    Args:
        results_df: DataFrame with experiment results
        metric: Metric to analyze
        split: Data split to show

    Returns:
        Plotly figure with violin plots

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    colors = get_palette("models", results_df["model"].nunique())
    color_map = {model: colors[i] for i, model in enumerate(results_df["model"].unique())}

    fig = px.violin(
        results_df,
        x="model",
        y=metric_col,
        color="model",
        box=True,
        points="all",
        title=f"{format_metric_name(metric)} Distribution by Model Type ({split.capitalize()} Set)",
        labels={metric_col: format_metric_name(metric), "model": "Model"},
        color_discrete_map=color_map,
    )

    fig.update_layout(
        height=500,
        showlegend=False,
    )

    return fig


def create_grid_vs_model_performance(
    results_df: pd.DataFrame,
    metric: str,
    split: str = "test",
) -> go.Figure:
    """Create a faceted plot showing model performance across grids.

    Args:
        results_df: DataFrame with experiment results
        metric: Metric to analyze
        split: Data split to show

    Returns:
        Plotly figure with faceted plots

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    results_df = results_df.copy()
    results_df["grid_level"] = results_df["grid"] + "_" + results_df["level"].astype(str)

    fig = px.box(
        results_df,
        x="model",
        y=metric_col,
        color="model",
        facet_col="grid_level",
        facet_row="task",
        points="all",
        title=f"{format_metric_name(metric)}: Model Performance Across Grid Levels and Tasks ({split.capitalize()})",
        labels={metric_col: format_metric_name(metric), "model": "Model"},
    )

    fig.update_layout(
        height=800,
        showlegend=True,
    )

    return fig


def create_performance_improvement_plot(
    results_df: pd.DataFrame,
    metric: str,
    baseline_model: str = "rf",
    split: str = "test",
) -> go.Figure:
    """Create a plot showing performance improvement relative to a baseline model.

    Args:
        results_df: DataFrame with experiment results
        metric: Metric to analyze
        baseline_model: Model to use as baseline for comparison
        split: Data split to show

    Returns:
        Plotly figure showing improvement over baseline

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in results_df.columns:
        raise ValueError(f"Metric {metric_col} not found in results")

    results_df = results_df.copy()

    # Calculate baseline performance for each task-target-grid-level combination
    baseline_perf = (
        results_df[results_df["model"] == baseline_model]
        .groupby(["task", "target", "grid", "level"])[metric_col]
        .max()
        .reset_index()
        .rename(columns={metric_col: "baseline_score"})
    )

    # Merge with all results
    results_with_baseline = results_df.merge(baseline_perf, on=["task", "target", "grid", "level"], how="left")

    # Calculate improvement
    results_with_baseline["improvement"] = results_with_baseline[metric_col] - results_with_baseline["baseline_score"]
    results_with_baseline["improvement_pct"] = (
        100 * results_with_baseline["improvement"] / results_with_baseline["baseline_score"]
    )

    # Filter out baseline model itself
    results_with_baseline = results_with_baseline[results_with_baseline["model"] != baseline_model]

    # Create plot
    fig = px.box(
        results_with_baseline,
        x="model",
        y="improvement_pct",
        color="model",
        points="all",
        title=(
            f"Performance Improvement Over {baseline_model.upper()} Baseline "
            f"({format_metric_name(metric)}, {split.capitalize()} Set)"
        ),
        labels={"improvement_pct": "Improvement (%)", "model": "Model"},
    )

    fig.add_hline(y=0, line_dash="dash", line_color="gray", annotation_text="Baseline")

    fig.update_layout(
        height=500,
        showlegend=False,
    )

    return fig


def create_top_models_bar_chart(
    task_df: pd.DataFrame,
    metric: str,
    task_name: str,
    split: str = "test",
    top_n: int = 10,
) -> go.Figure:
    """Create a horizontal bar chart showing top models for a task.

    Args:
        task_df: DataFrame filtered for a specific task
        metric: Metric to display
        task_name: Name of the task for the title
        split: Data split to show
        top_n: Number of top models to show

    Returns:
        Plotly figure with horizontal bar chart

    """
    metric_col = f"{split}_{metric}"

    if metric_col not in task_df.columns:
        raise ValueError(f"Metric {metric_col} not found in data")

    # Get top N models
    top_models = task_df.nlargest(top_n, metric_col).copy()

    # Create label combining model, grid level, and target
    top_models["label"] = (
        top_models["model"].str.upper()
        + " ("
        + top_models["grid"]
        + "-"
        + top_models["level"].astype(str)
        + ", "
        + top_models["target"]
        + ")"
    )

    # Sort by score for display (ascending so best is on top)
    top_models = top_models.sort_values(metric_col, ascending=True)

    # Create color map based on model type
    unique_models = top_models["model"].unique()
    model_colors = {model: get_palette("models", len(unique_models))[i] for i, model in enumerate(unique_models)}
    top_models["color"] = top_models["model"].map(model_colors)

    fig = go.Figure(
        data=[
            go.Bar(
                y=top_models["label"],
                x=top_models[metric_col],
                orientation="h",
                marker={"color": top_models["color"]},
                text=top_models[metric_col].round(4),
                textposition="auto",
                hovertemplate=("<b>%{y}</b><br>" + f"{format_metric_name(metric)}: %{{x:.4f}}<br>" + "<extra></extra>"),
            )
        ]
    )

    fig.update_layout(
        title=f"Top {top_n} Models - {task_name.replace('_', ' ').title()}",
        xaxis_title=format_metric_name(metric),
        yaxis_title="",
        height=max(400, top_n * 40),
        showlegend=False,
        margin={"l": 250},
    )

    return fig


def create_feature_importance_by_grid_level(
    fi_df: pd.DataFrame,
    top_n: int = 15,
) -> go.Figure:
    """Create a grouped bar chart showing top features by grid level.

    Args:
        fi_df: Feature importance DataFrame with columns: feature, importance, grid_level
        top_n: Number of top features to show

    Returns:
        Plotly figure with grouped bar chart

    """
    # Get overall top features
    overall_top = fi_df.groupby("feature")["importance"].mean().nlargest(top_n).index.tolist()

    # Filter to top features
    filtered = fi_df[fi_df["feature"].isin(overall_top)]

    # Calculate mean importance per feature per grid level
    grouped = filtered.groupby(["grid_level", "feature"])["importance"].mean().reset_index()

    # Create the plot
    fig = px.bar(
        grouped,
        x="feature",
        y="importance",
        color="grid_level",
        barmode="group",
        title=f"Top {top_n} Features by Grid Level",
        labels={"importance": "Mean Importance", "feature": "Feature", "grid_level": "Grid Level"},
    )

    fig.update_layout(
        height=500,
        xaxis_tickangle=-45,
        showlegend=True,
        legend={"orientation": "v", "yanchor": "top", "y": 1, "xanchor": "left", "x": 1.02},
    )

    return fig


def create_feature_consistency_plot(
    fi_df: pd.DataFrame,
    top_n: int = 15,
) -> go.Figure:
    """Create a scatter plot showing feature importance vs consistency.

    Args:
        fi_df: Feature importance DataFrame
        top_n: Number of top features to show

    Returns:
        Plotly figure with scatter plot

    """
    # Get top features
    overall_top = fi_df.groupby("feature")["importance"].mean().nlargest(top_n).index.tolist()

    # Calculate statistics
    stats = (
        fi_df[fi_df["feature"].isin(overall_top)].groupby("feature")["importance"].agg(["mean", "std"]).reset_index()
    )
    stats["cv"] = stats["std"] / stats["mean"]

    # Add data source if available
    if "data_source" in fi_df.columns:
        data_source_map = fi_df[["feature", "data_source"]].drop_duplicates().set_index("feature")["data_source"]
        stats["data_source"] = stats["feature"].map(data_source_map)
        color_col = "data_source"
    else:
        color_col = None

    fig = px.scatter(
        stats,
        x="cv",
        y="mean",
        size="std",
        color=color_col,
        hover_data=["feature"],
        text="feature",
        title="Feature Importance vs Consistency (Coefficient of Variation)",
        labels={
            "cv": "Coefficient of Variation (lower = more consistent)",
            "mean": "Mean Importance",
            "std": "Std Dev",
        },
    )

    fig.update_traces(textposition="top center", textfont_size=8)

    fig.update_layout(
        height=600,
        showlegend=True,
    )

    return fig


def create_data_source_importance_bars(
    fi_df: pd.DataFrame,
) -> go.Figure:
    """Create a bar chart showing importance breakdown by data source.

    Args:
        fi_df: Feature importance DataFrame with data_source column

    Returns:
        Plotly figure with bar chart

    """
    if "data_source" not in fi_df.columns:
        raise ValueError("data_source column not found in feature importance data")

    # Aggregate by data source
    source_stats = (
        fi_df.groupby("data_source")["importance"]
        .agg(["sum", "mean", "count"])
        .reset_index()
        .sort_values("sum", ascending=False)
    )

    # Create colors for data sources
    colors = get_palette("sources", len(source_stats))

    fig = go.Figure(
        data=[
            go.Bar(
                x=source_stats["data_source"],
                y=source_stats["sum"],
                marker={"color": colors},
                text=source_stats["sum"].round(2),
                textposition="auto",
                hovertemplate=(
                    "<b>%{x}</b><br>"
                    "Total Importance: %{y:.2f}<br>"
                    "Mean: %{customdata[0]:.4f}<br>"
                    "Features: %{customdata[1]}<br>"
                    "<extra></extra>"
                ),
                customdata=source_stats[["mean", "count"]].values,
            )
        ]
    )

    fig.update_layout(
        title="Feature Importance by Data Source",
        xaxis_title="Data Source",
        yaxis_title="Total Importance",
        height=400,
        showlegend=False,
    )

    return fig


def create_inference_maps(
    inference_gdf: "gpd.GeoDataFrame",
    grid: str,
    level: int,
    task: str,
) -> tuple["pdk.Deck", "pdk.Deck"]:
    """Create inference maps showing mean and std of predictions using pydeck.

    Args:
        inference_gdf: GeoDataFrame with geometry, mean_prediction, std_prediction columns
        grid: Grid type (e.g., 'hex', 'healpix')
        level: Grid resolution level
        task: Task type for title

    Returns:
        Tuple of (mean_deck, std_deck) pydeck visualizations

    """
    import matplotlib.colors as mcolors
    import numpy as np
    import pydeck as pdk
    from matplotlib.colors import LinearSegmentedColormap

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

    # Create a copy and convert to EPSG:4326 for pydeck
    gdf = inference_gdf.copy().to_crs("EPSG:4326")

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

    # Create custom colormap for predictions (white -> blue -> purple)
    colors_mean = ["#f7fbff", "#deebf7", "#c6dbef", "#9ecae1", "#6baed6", "#4292c6", "#2171b5", "#08519c", "#08306b"]
    cmap_mean = LinearSegmentedColormap.from_list("prediction", colors_mean)

    colors_std = ["#f7fcf5", "#e5f5e0", "#c7e9c0", "#a1d99b", "#74c476", "#41ab5d", "#238b45", "#006d2c", "#00441b"]
    cmap_std = LinearSegmentedColormap.from_list("uncertainty", colors_std)

    # Normalize mean predictions
    mean_values = gdf["mean_prediction"].to_numpy()
    mean_min, mean_max = np.nanpercentile(mean_values, [2, 98])
    if mean_max > mean_min:
        mean_normalized = np.clip((mean_values - mean_min) / (mean_max - mean_min), 0, 1)
    else:
        mean_normalized = np.zeros_like(mean_values)

    # Map normalized values to colors for mean
    mean_colors = [cmap_mean(val) for val in mean_normalized]
    mean_rgb_colors = [hex_to_rgb(mcolors.to_hex(color)) for color in mean_colors]
    gdf["mean_color"] = mean_rgb_colors

    # Normalize std predictions
    std_values = gdf["std_prediction"].to_numpy()
    std_min, std_max = np.nanpercentile(std_values, [2, 98])
    if std_max > std_min:
        std_normalized = np.clip((std_values - std_min) / (std_max - std_min), 0, 1)
    else:
        std_normalized = np.zeros_like(std_values)

    # Map normalized values to colors for std
    std_colors = [cmap_std(val) for val in std_normalized]
    std_rgb_colors = [hex_to_rgb(mcolors.to_hex(color)) for color in std_colors]
    gdf["std_color"] = std_rgb_colors

    # Convert to GeoJSON for mean predictions
    geojson_mean = []
    for _, row in gdf.iterrows():
        feature = {
            "type": "Feature",
            "geometry": row["geometry"].__geo_interface__,
            "properties": {
                "fill_color": row["mean_color"],
                "mean_prediction": float(row["mean_prediction"]),
            },
        }
        geojson_mean.append(feature)

    # Convert to GeoJSON for std predictions
    geojson_std = []
    for _, row in gdf.iterrows():
        feature = {
            "type": "Feature",
            "geometry": row["geometry"].__geo_interface__,
            "properties": {
                "fill_color": row["std_color"],
                "std_prediction": float(row["std_prediction"]),
            },
        }
        geojson_std.append(feature)

    # Create pydeck layer for mean predictions
    layer_mean = pdk.Layer(
        "GeoJsonLayer",
        geojson_mean,
        opacity=0.85,
        stroked=True,
        filled=True,
        extruded=False,
        wireframe=False,
        get_fill_color="properties.fill_color",
        get_line_color=[100, 100, 100],
        line_width_min_pixels=0.3,
        pickable=True,
    )

    # Create pydeck layer for std predictions
    layer_std = pdk.Layer(
        "GeoJsonLayer",
        geojson_std,
        opacity=0.85,
        stroked=True,
        filled=True,
        extruded=False,
        wireframe=False,
        get_fill_color="properties.fill_color",
        get_line_color=[100, 100, 100],
        line_width_min_pixels=0.3,
        pickable=True,
    )

    # Set initial view state for Arctic region
    view_state = pdk.ViewState(
        latitude=75,
        longitude=0,
        zoom=2.5,
        pitch=0,
    )

    # Create deck for mean predictions
    deck_mean = pdk.Deck(
        layers=[layer_mean],
        initial_view_state=view_state,
        tooltip={
            "html": (
                f"<b>Mean Prediction:</b> {{mean_prediction}}<br>"
                f"<b>{grid.upper()}-{level}</b> | <b>Task:</b> {task.title()}"
            ),
            "style": {"backgroundColor": "#2171b5", "color": "white"},
        },
        map_style="https://basemaps.cartocdn.com/gl/positron-gl-style/style.json",
    )

    # Create deck for std predictions
    deck_std = pdk.Deck(
        layers=[layer_std],
        initial_view_state=view_state,
        tooltip={
            "html": (
                f"<b>Uncertainty (Std):</b> {{std_prediction}}<br>"
                f"<b>{grid.upper()}-{level}</b> | <b>Task:</b> {task.title()}"
            ),
            "style": {"backgroundColor": "#238b45", "color": "white"},
        },
        map_style="https://basemaps.cartocdn.com/gl/positron-gl-style/style.json",
    )

    return deck_mean, deck_std
