"""Plots for visualizing target labels in datasets."""

import antimeridian
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import pydeck as pdk
from plotly.subplots import make_subplots
from shapely.geometry import shape

from entropice.dashboard.utils.colors import get_cmap, get_palette, hex_to_rgb
from entropice.ml.dataset import TrainingSet
from entropice.utils.types import TargetDataset, Task


def create_target_distribution_plot(train_data_dict: dict[TargetDataset, dict[Task, TrainingSet]]) -> go.Figure:
    """Create a plot showing the distribution of target labels across datasets and tasks.

    This function creates a comprehensive visualization showing:
    - Classification tasks (binary, count_regimes, density_regimes): categorical distributions
    - Regression tasks (count, density): histogram distributions
    - Train/test split information for each combination
    - Separate subplots for each TargetDataset

    Args:
        train_data_dict: Nested dictionary with structure:
            {target_dataset: {task: TrainingSet}}
            where target_dataset is in ["darts_v1", "darts_mllabels"]
            and task is in ["binary", "count_regimes", "density_regimes", "count", "density"]

    Returns:
        Plotly Figure with subplots showing target distributions.

    """
    # Define task types and their properties
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]

    task_titles: dict[Task, str] = {
        "binary": "Binary",
        "count_regimes": "Count Regimes",
        "density_regimes": "Density Regimes",
        "count": "Count (Regression)",
        "density": "Density (Regression)",
    }

    # Get all available target datasets and tasks
    target_datasets = sorted(train_data_dict.keys())
    all_tasks = sorted(
        {task for tasks in train_data_dict.values() for task in tasks.keys()},
        key=lambda x: (x not in classification_tasks, x),  # Classification first
    )

    # Create subplots: one row per target dataset, one column per task
    n_rows = len(target_datasets)
    n_cols = len(all_tasks)

    # Create column titles (tasks) for the subplots
    column_titles = [task_titles.get(task, str(task)) for task in all_tasks]  # type: ignore[arg-type]

    # Create row titles (target datasets)
    row_titles = [target_ds.replace("_", " ").title() for target_ds in target_datasets]

    fig = make_subplots(
        rows=n_rows,
        cols=n_cols,
        column_titles=column_titles,
        vertical_spacing=0.20 / max(n_rows, 1),
        horizontal_spacing=0.08 / max(n_cols, 1),
    )

    # Iterate through each target dataset and task
    for row_idx, target_dataset in enumerate(target_datasets, start=1):
        task_dict = train_data_dict[target_dataset]  # type: ignore[index]

        for col_idx, task in enumerate(all_tasks, start=1):
            if task not in task_dict:
                # Skip if this task is not available for this target dataset
                continue

            dataset = task_dict[task]  # type: ignore[index]

            # Determine if this is a classification or regression task
            if task in classification_tasks:
                _add_classification_subplot(fig, dataset, task, row_idx, col_idx, n_rows)
            else:
                _add_regression_subplot(fig, dataset, task, row_idx, col_idx, n_rows)

    # Update layout
    fig.update_layout(
        height=400 * n_rows,
        showlegend=True,
        legend={
            "orientation": "h",
            "yanchor": "top",
            "y": -0.15,
            "xanchor": "center",
            "x": 0.5,
        },
        margin={"l": 120, "r": 40, "t": 80, "b": 100},
    )

    # Update all x-axes and y-axes
    fig.update_xaxes(title_font={"size": 11})
    fig.update_yaxes(title_text="Count", title_font={"size": 11})

    # Manually add row title annotations (left-aligned dataset names)
    for row_idx, row_title in enumerate(row_titles, start=1):
        # Calculate the y position for each row (at the top of the row)
        # Account for vertical spacing between rows
        spacing = 0.20 / max(n_rows, 1)
        row_height = (1 - spacing * (n_rows - 1)) / n_rows
        y_position = 1.08 - (row_idx - 1) * (row_height + spacing)

        fig.add_annotation(
            text=f"<b>{row_title}</b>",
            xref="paper",
            yref="paper",
            x=-0.02,
            y=y_position,
            xanchor="left",
            yanchor="top",
            textangle=0,
            font={"size": 20},
            showarrow=False,
        )

    return fig


def _add_classification_subplot(
    fig: go.Figure,
    dataset: TrainingSet,
    task: str,
    row: int,
    col: int,
    n_rows: int,
) -> None:
    """Add a classification task subplot showing categorical distribution.

    Args:
        fig: Plotly figure to add subplot to.
        dataset: TrainingSet containing the data.
        task: Task name for color selection.
        row: Subplot row index.
        col: Subplot column index.
        n_rows: Total number of rows (for legend control).

    """
    y_binned = dataset.targets["y"]
    categories = y_binned.cat.categories.tolist()
    colors = get_palette(task, len(categories) + 2)[1:-1]  # Avoid too light/dark colors

    # Calculate counts for train and test splits
    train_counts = [((y_binned == cat) & (dataset.split == "train")).sum() for cat in categories]
    test_counts = [((y_binned == cat) & (dataset.split == "test")).sum() for cat in categories]

    # Add train bars
    fig.add_trace(
        go.Bar(
            name="Train",
            x=categories,
            y=train_counts,
            marker_color=colors,
            opacity=0.9,
            text=train_counts,
            textposition="inside",
            textfont={"size": 9, "color": "white"},
            legendgroup="split",
            showlegend=(row == 1 and col == 1),  # Only show legend once
        ),
        row=row,
        col=col,
    )

    # Add test bars
    fig.add_trace(
        go.Bar(
            name="Test",
            x=categories,
            y=test_counts,
            marker_color=colors,
            opacity=0.5,
            text=test_counts,
            textposition="inside",
            textfont={"size": 9, "color": "white"},
            legendgroup="split",
            showlegend=(row == 1 and col == 1),  # Only show legend once
        ),
        row=row,
        col=col,
    )

    # Update subplot layout
    fig.update_xaxes(tickangle=-45, row=row, col=col)


def _add_regression_subplot(
    fig: go.Figure,
    dataset: TrainingSet,
    task: str,
    row: int,
    col: int,
    n_rows: int,
) -> None:
    """Add a regression task subplot showing histogram distribution.

    Args:
        fig: Plotly figure to add subplot to.
        dataset: TrainingSet containing the data.
        task: Task name for color selection.
        row: Subplot row index.
        col: Subplot column index.
        n_rows: Total number of rows (for legend control).

    """
    # Get raw values
    z_values = dataset.targets["z"]

    # Split into train and test
    train_values = z_values[dataset.split == "train"]
    test_values = z_values[dataset.split == "test"]

    # Determine bin edges based on complete data
    all_values = pd.concat([train_values, test_values])
    # Use a reasonable number of bins
    n_bins = min(30, int(np.sqrt(len(all_values))))
    bin_edges = np.histogram_bin_edges(all_values, bins=n_bins)

    # Get colors for the task
    colors = get_palette(task, 6)[1:5:2]  # Use less bright/dark colors

    # Add train histogram
    fig.add_trace(
        go.Histogram(
            name="Train",
            x=train_values,
            xbins={"start": bin_edges[0], "end": bin_edges[-1], "size": bin_edges[1] - bin_edges[0]},
            marker_color=colors[0],
            opacity=0.7,
            legendgroup="split",
            showlegend=(row == 1 and col == 1),  # Only show legend once
        ),
        row=row,
        col=col,
    )

    # Add test histogram
    fig.add_trace(
        go.Histogram(
            name="Test",
            x=test_values,
            xbins={"start": bin_edges[0], "end": bin_edges[-1], "size": bin_edges[1] - bin_edges[0]},
            marker_color=colors[1],
            opacity=0.5,
            legendgroup="split",
            showlegend=(row == 1 and col == 1),  # Only show legend once
        ),
        row=row,
        col=col,
    )

    # Update barmode to overlay for histograms
    fig.update_xaxes(title_text=task.capitalize(), row=row, col=col)


def _assign_split_colors(gdf: pd.DataFrame) -> pd.DataFrame:
    """Assign colors based on train/test split.

    Args:
        gdf: GeoDataFrame with 'split' column.

    Returns:
        GeoDataFrame with 'fill_color' column added.

    """
    split_colors = get_palette("split", 2)
    color_map = {"train": hex_to_rgb(split_colors[0]), "test": hex_to_rgb(split_colors[1])}
    # Convert to list to avoid categorical hashing issues with list values
    gdf["fill_color"] = [color_map[x] for x in gdf["split"]]
    return gdf


def _assign_classification_colors(gdf: pd.DataFrame, task: Task) -> pd.DataFrame:
    """Assign colors for classification tasks based on categorical labels.

    Args:
        gdf: GeoDataFrame with categorical 'y' column.
        task: Task name for color selection.

    Returns:
        GeoDataFrame with 'fill_color' column added.

    """
    categories = gdf["y"].cat.categories.tolist()
    colors_palette = get_palette(task, len(categories))
    color_map = {cat: hex_to_rgb(colors_palette[i]) for i, cat in enumerate(categories)}
    # Convert to list to avoid categorical hashing issues with list values
    gdf["fill_color"] = [color_map[x] for x in gdf["y"]]
    return gdf


def _assign_regression_colors(gdf: pd.DataFrame, task: Task) -> pd.DataFrame:
    """Assign colors for regression tasks based on continuous values.

    Args:
        gdf: GeoDataFrame with numeric 'z' column.
        task: Task name for color selection.

    Returns:
        GeoDataFrame with 'fill_color' column added.

    """
    z_values = gdf["z"]
    z_min, z_max = z_values.min(), z_values.max()

    # Normalize to [0, 1]
    if z_max > z_min:
        z_normalized = (z_values - z_min) / (z_max - z_min)
    else:
        z_normalized = pd.Series([0.5] * len(z_values), index=z_values.index)

    # Get colormap and map normalized values to RGB
    cmap = get_cmap(task)

    def value_to_rgb(normalized_val):
        rgba = cmap(normalized_val)
        return [int(rgba[0] * 255), int(rgba[1] * 255), int(rgba[2] * 255)]

    gdf["fill_color"] = z_normalized.apply(value_to_rgb)
    return gdf


def _assign_elevation(gdf: pd.DataFrame, make_3d_map: bool, is_classification: bool) -> pd.DataFrame:
    """Assign elevation values for 3D visualization.

    Args:
        gdf: GeoDataFrame to add elevation to.
        make_3d_map: Whether to create a 3D map.
        is_classification: Whether this is a classification task.

    Returns:
        GeoDataFrame with 'elevation' column added.

    """
    if not make_3d_map:
        gdf["elevation"] = 0
        return gdf

    print(f"{make_3d_map=}, {is_classification=}")
    if is_classification:
        # For classification, use category codes for elevation
        # Convert to int to avoid overflow with int8 categorical codes
        z_values = gdf["y"].cat.codes.astype(int)
    else:
        # For regression, use z values for elevation
        z_values = gdf["z"]
    z_min, z_max = z_values.min(), z_values.max()

    if z_max > z_min:
        gdf["elevation"] = ((z_values - z_min) / (z_max - z_min)).fillna(0)
    else:
        gdf["elevation"] = 0

    print(gdf["elevation"].describe())
    return gdf


def create_target_spatial_distribution_map(
    training_set: TrainingSet, make_3d_map: bool, show_split: bool, task: Task
) -> pdk.Deck:
    """Create a spatial distribution map of target labels using PyDeck.

    Args:
        training_set (TrainingSet): The training set containing spatial data and target labels.
        make_3d_map (bool): Whether to create a 3D map (True) or a 2D map (False).
        show_split (bool): Whether to color by train/test split (True) or by target labels (False).
        task (Task): The task type to determine color mapping.

    Returns:
        pdk.Deck: A PyDeck Deck object representing the spatial distribution map.

    """
    # Get the targets GeoDataFrame
    gdf = training_set.targets.copy()

    # Add split information
    gdf["split"] = training_set.split.to_numpy()

    # Determine if this is a classification or regression task
    is_classification = gdf["y"].dtype.name == "category"

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

    # Fix antimeridian issues
    def fix_hex_geometry(geom):
        try:
            return shape(antimeridian.fix_shape(geom))
        except (ValueError, Exception):
            return geom

    gdf_wgs84["geometry"] = gdf_wgs84["geometry"].apply(fix_hex_geometry)

    # Assign colors based on mode
    if show_split:
        gdf_wgs84 = _assign_split_colors(gdf_wgs84)
    elif is_classification:
        gdf_wgs84 = _assign_classification_colors(gdf_wgs84, task)
    else:
        gdf_wgs84 = _assign_regression_colors(gdf_wgs84, task)

    # Add elevation for 3D visualization
    gdf_wgs84 = _assign_elevation(gdf_wgs84, make_3d_map, is_classification)

    # 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"],
                "elevation": row["elevation"],
                "split": row["split"],
                "y_label": str(row["y"]),
                "z_value": float(row["z"]) if pd.notna(row["z"]) else None,
            },
        }
        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 if make_3d_map else 0,
        pickable=True,
    )

    # Set initial view state (centered on the Arctic)
    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,
    )

    # Create tooltip
    tooltip_html = "<b>Split:</b> {split}<br/><b>Label:</b> {y_label}<br/><b>Value:</b> {z_value}<br/>"

    # 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
