"""Visualization functions for the overview page.

This module contains reusable plotting functions for dataset analysis visualizations,
including sample counts, feature counts, and dataset statistics.
"""

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from entropice.dashboard.utils.colors import get_palette


def _get_colors_from_members(unique_members: list[str]) -> dict[str, str]:
    unique_members = sorted(unique_members)
    unique_sources = {member.split("-")[0] for member in unique_members}
    source_member_map = {
        source: [member for member in unique_members if member.split("-")[0] == source] for source in unique_sources
    }
    source_color_map = {}
    for source, members in source_member_map.items():
        n_members = len(members) + 2  # Avoid very light/dark colors
        palette = get_palette(source, n_colors=n_members)
        for idx, member in enumerate(members, start=1):
            source_color_map[member] = palette[idx]
    return source_color_map


def create_sample_count_bar_chart(sample_df: pd.DataFrame) -> go.Figure:
    """Create bar chart showing sample counts by grid, target and temporal mode.

    Args:
        sample_df: DataFrame with columns: Grid, Target, Samples (Coverage).

    Returns:
        Plotly Figure object containing the bar chart visualization.

    """
    assert "Grid" in sample_df.columns
    assert "Target" in sample_df.columns
    assert "Temporal Mode" in sample_df.columns
    assert "Samples (Coverage)" in sample_df.columns

    targets = sorted(sample_df["Target"].unique())
    modes = sorted(sample_df["Temporal Mode"].unique())

    # Create subplots manually to have better control over colors
    fig = make_subplots(
        rows=1,
        cols=len(targets),
        subplot_titles=[f"Target: {target}" for target in targets],
        shared_yaxes=True,
    )
    # Get color palette for this plot
    colors = get_palette("Number of Samples", n_colors=len(modes))

    # Track which modes have been added to legend
    modes_in_legend = set()

    for col_idx, target in enumerate(targets, start=1):
        target_data = sample_df[sample_df["Target"] == target]

        for mode_idx, mode in enumerate(modes):
            mode_data = target_data[target_data["Temporal Mode"] == mode]
            if len(mode_data) == 0:
                continue
            color = colors[mode_idx] if colors and mode_idx < len(colors) else None

            # Show legend only the first time each mode appears
            show_in_legend = mode not in modes_in_legend
            if show_in_legend:
                modes_in_legend.add(mode)

            # Create a unique legendgroup per mode so colors are consistent
            fig.add_trace(
                go.Bar(
                    x=mode_data["Grid"],
                    y=mode_data["Samples (Coverage)"],
                    name=mode,
                    marker_color=color,
                    legendgroup=mode,  # Group by mode name
                    showlegend=show_in_legend,
                ),
                row=1,
                col=col_idx,
            )

    fig.update_layout(
        title_text="Training Sample Counts by Grid Configuration and Target Dataset",
        barmode="group",
        height=500,
        showlegend=True,
    )

    fig.update_xaxes(tickangle=-45)
    fig.update_yaxes(title_text="Number of Samples", row=1, col=1)

    return fig


def create_feature_count_stacked_bar(breakdown_df: pd.DataFrame) -> go.Figure:
    """Create stacked bar chart showing feature counts by member.

    Args:
        breakdown_df: DataFrame with columns: Grid, Temporal Mode, Member, Number of Features.

    Returns:
        Plotly Figure object containing the stacked bar chart visualization.

    """
    assert "Grid" in breakdown_df.columns
    assert "Temporal Mode" in breakdown_df.columns
    assert "Member" in breakdown_df.columns
    assert "Number of Features" in breakdown_df.columns

    unique_members = sorted(breakdown_df["Member"].unique())
    temporal_modes = sorted(breakdown_df["Temporal Mode"].unique())
    source_color_map = _get_colors_from_members(unique_members)

    # Create subplots for each temporal mode
    fig = make_subplots(
        rows=1,
        cols=len(temporal_modes),
        subplot_titles=[f"Temporal Mode: {mode}" for mode in temporal_modes],
        shared_yaxes=True,
    )

    for col_idx, mode in enumerate(temporal_modes, start=1):
        mode_data = breakdown_df[breakdown_df["Temporal Mode"] == mode]

        for member in unique_members:
            member_data = mode_data[mode_data["Member"] == member]
            color = source_color_map.get(member)

            fig.add_trace(
                go.Bar(
                    x=member_data["Grid"],
                    y=member_data["Number of Features"],
                    name=member,
                    marker_color=color,
                    legendgroup=member,
                    showlegend=(col_idx == 1),
                ),
                row=1,
                col=col_idx,
            )

    fig.update_layout(
        title_text="Input Features by Member Across Grid Configurations",
        barmode="stack",
        height=500,
        showlegend=True,
    )
    fig.update_xaxes(tickangle=-45)
    fig.update_yaxes(title_text="Number of Features", row=1, col=1)

    return fig


def create_inference_cells_bar(sample_df: pd.DataFrame) -> go.Figure:
    """Create bar chart for inference cells by grid configuration.

    Args:
        sample_df: DataFrame with columns: Grid, Temporal Mode, Inference Cells.

    Returns:
        Plotly Figure object containing the bar chart visualization.

    """
    assert "Grid" in sample_df.columns
    assert "Inference Cells" in sample_df.columns

    n_grids = sample_df["Grid"].nunique()
    mode_colors: list[str] = get_palette("Number of Samples", n_colors=n_grids)

    fig = px.bar(
        sample_df,
        x="Grid",
        y="Inference Cells",
        color="Grid",
        title="Spatial Coverage (Grid Cells with Complete Data)",
        labels={
            "Grid": "Grid Configuration",
            "Inference Cells": "Number of Cells",
        },
        color_discrete_sequence=mode_colors,
        text="Inference Cells",
    )

    fig.update_traces(texttemplate="%{text:,}", textposition="outside")
    fig.update_layout(xaxis_tickangle=-45, height=500, showlegend=False)

    return fig


def create_total_samples_bar(
    comparison_df: pd.DataFrame,
    grid_colors: list[str] | None = None,
) -> go.Figure:
    """Create bar chart for total samples by grid configuration.

    Args:
        comparison_df: DataFrame with columns: Grid, Total Samples.
        grid_colors: Optional color palette for grid configurations. If None, uses default Plotly colors.

    Returns:
        Plotly Figure object containing the bar chart visualization.

    """
    fig = px.bar(
        comparison_df,
        x="Grid",
        y="Total Samples",
        color="Grid",
        title="Training Samples (Binary Task)",
        labels={
            "Grid": "Grid Configuration",
            "Total Samples": "Number of Samples",
        },
        color_discrete_sequence=grid_colors,
        text="Total Samples",
    )

    fig.update_traces(texttemplate="%{text:,}", textposition="outside")
    fig.update_layout(xaxis_tickangle=-45, showlegend=False)

    return fig


def create_feature_breakdown_donut(
    grid_data: pd.DataFrame,
    grid_config: str,
    source_color_map: dict[str, str] | None = None,
) -> go.Figure:
    """Create donut chart for feature breakdown by data source for a specific grid.

    Args:
        grid_data: DataFrame with columns: Data Source, Number of Features.
        grid_config: Grid configuration name for the title.
        source_color_map: Optional dictionary mapping data source names to specific colors.
            If None, uses default Plotly colors.

    Returns:
        Plotly Figure object containing the donut chart visualization.

    """
    fig = px.pie(
        grid_data,
        names="Data Source",
        values="Number of Features",
        title=grid_config,
        hole=0.4,
        color_discrete_map=source_color_map,
        color="Data Source",
    )

    fig.update_traces(textposition="inside", textinfo="percent")
    fig.update_layout(showlegend=True, height=350)

    return fig


def create_feature_distribution_pie(
    breakdown_df: pd.DataFrame,
    source_color_map: dict[str, str] | None = None,
) -> go.Figure:
    """Create pie chart for feature distribution by data source.

    Args:
        breakdown_df: DataFrame with columns: Data Source, Number of Features.
        source_color_map: Optional dictionary mapping data source names to specific colors.
            If None, uses default Plotly colors.

    Returns:
        Plotly Figure object containing the pie chart visualization.

    """
    fig = px.pie(
        breakdown_df,
        names="Data Source",
        values="Number of Features",
        title="Feature Distribution by Data Source",
        hole=0.4,
        color_discrete_map=source_color_map,
        color="Data Source",
    )

    fig.update_traces(textposition="inside", textinfo="percent+label")
    fig.update_layout(height=400)

    return fig


def create_member_details_table(member_stats_dict: dict[str, dict]) -> pd.DataFrame:
    """Create a detailed table for member statistics.

    Args:
        member_stats_dict: Dictionary mapping member names to their statistics.
            Each stats dict should contain: feature_count, variable_names, dimensions, size_bytes.

    Returns:
        DataFrame with detailed member information.

    """
    rows = []
    for member, stats in member_stats_dict.items():
        # Calculate total data points
        total_points = 1
        for dim_size in stats["dimensions"].values():
            total_points *= dim_size

        # Format dimension string
        dim_str = " x ".join([f"{k}={v:,}" for k, v in stats["dimensions"].items()])

        rows.append(
            {
                "Member": member,
                "Features": stats["feature_count"],
                "Variables": len(stats["variable_names"]),
                "Dimensions": dim_str,
                "Data Points": f"{total_points:,}",
                "Size (MB)": f"{stats['size_bytes'] / (1024 * 1024):.2f}",
            }
        )

    return pd.DataFrame(rows)


def create_feature_breakdown_donuts_grid(breakdown_df: pd.DataFrame) -> go.Figure:
    """Create grid of donut charts showing feature breakdown by member for each grid+temporal mode.

    Args:
        breakdown_df: DataFrame with columns: Grid, Temporal Mode, Member, Number of Features.

    Returns:
        Plotly Figure object containing the grid of donut charts.

    """
    assert "Grid" in breakdown_df.columns
    assert "Temporal Mode" in breakdown_df.columns
    assert "Member" in breakdown_df.columns
    assert "Number of Features" in breakdown_df.columns

    # Get unique grids and temporal modes
    grids = breakdown_df["Grid"].unique().tolist()
    temporal_modes = breakdown_df["Temporal Mode"].unique().tolist()
    unique_members = sorted(breakdown_df["Member"].unique())
    source_color_map = _get_colors_from_members(unique_members)

    n_grids = len(grids)
    n_modes = len(temporal_modes)

    # Create subplot grid (grids as rows, temporal modes as columns)
    # Don't use subplot_titles - we'll add custom annotations for row/column labels
    fig = make_subplots(
        rows=n_grids,
        cols=n_modes,
        specs=[[{"type": "pie"}] * n_modes for _ in range(n_grids)],
        horizontal_spacing=0.01,
        vertical_spacing=0.02,
    )

    # Get all members to ensure consistent coloring across all donuts
    all_members = sorted(breakdown_df["Member"].unique())

    for row_idx, grid in enumerate(grids, start=1):
        for col_idx, mode in enumerate(temporal_modes, start=1):
            subset = breakdown_df[(breakdown_df["Grid"] == grid) & (breakdown_df["Temporal Mode"] == mode)]

            if len(subset) == 0:
                continue

            # Build labels, values, and colors in consistent order
            labels = []
            values = []
            colors_list = []

            for member in all_members:
                member_data = subset[subset["Member"] == member]
                if len(member_data) > 0:
                    labels.append(member)
                    values.append(member_data["Number of Features"].iloc[0])
                    if source_color_map:
                        colors_list.append(source_color_map[member])

            # Only show legend for first subplot
            show_legend = row_idx == 1 and col_idx == 1

            fig.add_trace(
                go.Pie(
                    labels=labels,
                    values=values,
                    name=f"{grid} - {mode}",
                    hole=0.4,
                    marker={"colors": colors_list} if colors_list else None,
                    textposition="inside",
                    textinfo="percent",
                    showlegend=show_legend,
                ),
                row=row_idx,
                col=col_idx,
            )

    # Calculate appropriate height based on number of rows
    height = max(500, n_grids * 400)

    fig.update_layout(
        title_text="Feature Breakdown by Member Across Grid Configurations and Temporal Modes",
        height=height,
        showlegend=True,
        margin={"l": 100, "r": 50, "t": 150, "b": 50},  # Add left margin for row labels
    )

    # Calculate subplot positions in paper coordinates
    # Account for spacing between subplots
    h_spacing = 0.01
    v_spacing = 0.02

    # Calculate width and height of each subplot in paper coordinates
    subplot_width = (1.0 - h_spacing * (n_modes - 1)) / n_modes
    subplot_height = (1.0 - v_spacing * (n_grids - 1)) / n_grids

    # Add column headers at the top (temporal modes)
    for col_idx, mode in enumerate(temporal_modes):
        # Calculate x position for this column's center
        x_pos = col_idx * (subplot_width + h_spacing) + subplot_width / 2
        fig.add_annotation(
            text=f"<b>{mode}</b>",
            xref="paper",
            yref="paper",
            x=x_pos,
            y=1.01,
            showarrow=False,
            xanchor="center",
            yanchor="bottom",
            font={"size": 18},
        )

    # Add row headers on the left (grid configurations)
    for row_idx, grid in enumerate(grids):
        # Calculate y position for this row's center
        # Note: y coordinates go from bottom to top, but rows go from top to bottom
        y_pos = 1.0 - (row_idx * (subplot_height + v_spacing) + subplot_height / 2)
        fig.add_annotation(
            text=f"<b>{grid}</b>",
            xref="paper",
            yref="paper",
            x=-0.01,
            y=y_pos,
            showarrow=False,
            xanchor="right",
            yanchor="middle",
            font={"size": 18},
            textangle=-90,
        )

    return fig
