"""Target dataset dashboard section."""

from typing import cast

import matplotlib.colors as mcolors
import streamlit as st

from entropice.dashboard.plots.targets import create_target_distribution_plot, create_target_spatial_distribution_map
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 _render_distribution(train_data_dict: dict[TargetDataset, dict[Task, TrainingSet]]):
    st.subheader("Target Labels Distribution")
    fig = create_target_distribution_plot(train_data_dict)
    st.plotly_chart(fig, width="stretch")
    st.markdown(
        """
        The above plot shows the distribution of target labels in the training and test datasets.
        You can use this information to understand the balance of classes and identify any potential
        issues with class imbalance that may affect model training.
        """
    )


def _render_split_legend(training_set: TrainingSet):
    """Render legend for train/test split visualization."""
    st.markdown("**Data Split:**")
    split_colors = get_palette("split", 2)

    col1, col2 = st.columns(2)
    with col1:
        color_rgb = hex_to_rgb(split_colors[0])
        st.markdown(
            f'<div style="display: flex; align-items: center;">'
            f'<div style="width: 20px; height: 20px; '
            f"background-color: rgb({color_rgb[0]}, {color_rgb[1]}, {color_rgb[2]}); "
            f'margin-right: 8px; border: 1px solid #ccc;"></div>'
            f"<span>Train</span></div>",
            unsafe_allow_html=True,
        )
    with col2:
        color_rgb = hex_to_rgb(split_colors[1])
        st.markdown(
            f'<div style="display: flex; align-items: center;">'
            f'<div style="width: 20px; height: 20px; '
            f"background-color: rgb({color_rgb[0]}, {color_rgb[1]}, {color_rgb[2]}); "
            f'margin-right: 8px; border: 1px solid #ccc;"></div>'
            f"<span>Test</span></div>",
            unsafe_allow_html=True,
        )

    # Show split statistics
    train_count = (training_set.split == "train").sum()
    test_count = (training_set.split == "test").sum()
    total = len(training_set)
    st.caption(
        f"Train: {train_count:,} ({train_count / total * 100:.1f}%) | "
        f"Test: {test_count:,} ({test_count / total * 100:.1f}%)"
    )


def _render_classification_legend(training_set: TrainingSet, selected_task: Task):
    """Render legend for classification task visualization."""
    st.markdown("**Target Classes:**")
    categories = training_set.targets["y"].cat.categories.tolist()
    colors_palette = get_palette(selected_task, len(categories))
    intervals = training_set.target_intervals

    # Show categories with colors and intervals
    for i, cat in enumerate(categories):
        color = colors_palette[i]
        interval_min, interval_max = intervals[i]

        # Format interval display
        interval_str = _format_interval(interval_min, interval_max, selected_task)

        st.markdown(
            f'<div style="display: flex; align-items: center; margin-bottom: 4px;">'
            f'<div style="width: 20px; height: 20px; background-color: {color}; '
            f'margin-right: 8px; border: 1px solid #ccc; flex-shrink: 0;"></div>'
            f"<span>{cat}{interval_str}</span></div>",
            unsafe_allow_html=True,
        )

    # Show total samples
    st.caption(f"Total samples: {len(training_set):,}")


def _format_interval(interval_min, interval_max, selected_task: Task) -> str:
    """Format interval string based on task type."""
    if interval_min is None or interval_max is None:
        return ""

    if selected_task in ["count", "count_regimes"]:
        # Integer values for count
        if interval_min == interval_max:
            return f" ({int(interval_min)})"
        return f" ({int(interval_min)}-{int(interval_max)})"

    if selected_task in ["density", "density_regimes"]:
        # Percentage values for density
        if interval_min == interval_max:
            return f" ({interval_min * 100:.4f}%)"
        return f" ({interval_min * 100:.4f}%-{interval_max * 100:.4f}%)"

    # Binary or other
    return ""


def _render_regression_legend(training_set: TrainingSet, selected_task: Task):
    """Render legend for regression task visualization."""
    st.markdown("**Target Values:**")
    z_values = training_set.targets["z"]
    z_min, z_max = z_values.min(), z_values.max()
    z_mean = z_values.mean()
    z_median = z_values.median()

    st.markdown(f"- **Range:** {z_min:.4f} to {z_max:.4f}\n- **Mean:** {z_mean:.4f}\n- **Median:** {z_median:.4f}")

    # Show gradient visualization using the actual colormap
    st.markdown("**Color Gradient:**")

    # Get the same colormap used in the map
    cmap = get_cmap(selected_task)
    # Sample colors from the colormap
    n_samples = 10
    gradient_colors = [mcolors.to_hex(cmap(i / (n_samples - 1))) for i in range(n_samples)]
    gradient_str = ", ".join(gradient_colors)

    st.markdown(
        f'<div style="display: flex; align-items: center; margin-top: 8px;">'
        f'<span style="margin-right: 8px;">Low ({z_min:.4f})</span>'
        f'<div style="width: 200px; height: 20px; '
        f"background: linear-gradient(to right, {gradient_str}); "
        f'border: 1px solid #ccc;"></div>'
        f'<span style="margin-left: 8px;">High ({z_max:.4f})</span>'
        f"</div>",
        unsafe_allow_html=True,
    )


def _render_elevation_info(training_set: TrainingSet):
    """Render 3D elevation information."""
    st.markdown("---")
    st.markdown("**3D Elevation:**")

    is_classification = training_set.targets["y"].dtype.name == "category"
    if is_classification:
        st.info("💡 Height represents category level. Rotate the map by holding Ctrl/Cmd and dragging with your mouse.")
    else:
        z_values = training_set.targets["z"]
        z_min, z_max = z_values.min(), z_values.max()
        st.info(
            f"💡 Height represents target value (normalized): {z_min:.4f} (low) → {z_max:.4f} (high). "
            "Rotate the map by holding Ctrl/Cmd and dragging with your mouse."
        )


@st.fragment
def _render_target_map(train_data_dict: dict[TargetDataset, dict[Task, TrainingSet]]):
    st.subheader("Target Labels Spatial Distribution")

    cols = st.columns([2, 2, 1])
    with cols[0]:
        selected_target = cast(
            TargetDataset,
            st.selectbox(
                "Select Target Dataset",
                options=sorted(train_data_dict.keys()),
                index=0,
                key="target_map_dataset",
            ),
        )
    with cols[1]:
        selected_task = cast(
            Task,
            st.selectbox(
                "Select Task",
                options=sorted(train_data_dict[selected_target].keys()),
                index=0,
                key="target_map_task",
            ),
        )
    with cols[2]:
        # Controls weather a 3D map or a 2D map is shown
        make_3d_map = cast(bool, st.toggle("3D Map", value=True, key="target_map_3d"))
        # Controls what should be shows, either the split or the labels / values
        show_split = cast(bool, st.checkbox("Show Train/Test Split", value=False, key="target_map_show_split"))

    training_set = train_data_dict[selected_target][selected_task]
    map_deck = create_target_spatial_distribution_map(training_set, make_3d_map, show_split, selected_task)
    st.pydeck_chart(map_deck)

    # Add legend
    with st.expander("🎨 Legend", expanded=True):
        if show_split:
            _render_split_legend(training_set)
        else:
            # Show target labels legend
            is_classification = training_set.targets["y"].dtype.name == "category"
            if is_classification:
                _render_classification_legend(training_set, selected_task)
            else:
                _render_regression_legend(training_set, selected_task)

        # Add elevation info for 3D maps
        if make_3d_map:
            _render_elevation_info(training_set)


def render_target_information_tab(train_data_dict: dict[TargetDataset, dict[Task, TrainingSet]]):
    """Render target labels distribution and spatial visualization.

    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"]

    """
    _render_distribution(train_data_dict)

    st.divider()

    _render_target_map(train_data_dict)
