"""Training Result Sections."""

import streamlit as st

from entropice.dashboard.plots.metrics import plot_confusion_matrix
from entropice.dashboard.utils.formatters import format_metric_name
from entropice.dashboard.utils.loaders import TrainingResult
from entropice.dashboard.utils.stats import CVMetricStatistics
from entropice.utils.types import GridConfig


def render_run_information(selected_result: TrainingResult, refit_metric):
    """Render training run configuration overview.

    Args:
        selected_result: The selected TrainingResult object.
        refit_metric: The refit metric used for model selection.

    """
    st.header("📋 Run Information")

    grid_config = GridConfig.from_grid_level((selected_result.run.dataset.grid, selected_result.run.dataset.level))

    col1, col2, col3, col4, col5 = st.columns(5)
    with col1:
        st.metric("Task", selected_result.run.task.capitalize())
    with col2:
        st.metric("Target", selected_result.run.target.capitalize())
    with col3:
        st.metric("Grid", grid_config.display_name)
    with col4:
        st.metric("Model", selected_result.run.model_type.upper())
    with col5:
        st.metric("Trials", selected_result.run.n_trials or "N/A")

    st.caption(f"**Refit Metric:** {format_metric_name(refit_metric)}")


def _render_metrics(metrics: dict[str, float]):
    """Render a set of metrics in a max-five-column layout.

    Args:
        metrics: Dictionary of metric names and their values.

    """
    ncols = min(5, len(metrics))
    cols = st.columns(ncols)
    for idx, (metric_name, metric_value) in enumerate(metrics.items()):
        with cols[idx % ncols]:
            st.metric(format_metric_name(metric_name), f"{metric_value:.4f}")


def render_metrics_section(selected_result: TrainingResult):
    """Render test metrics overview showing final model performance.

    Args:
        selected_result: The selected TrainingResult object.

    """
    # Extract metrics for each split
    test_metrics = selected_result.run.get_metrics_from_split("test")
    train_metrics = selected_result.run.get_metrics_from_split("train")
    complete_metrics = selected_result.run.get_metrics_from_split("complete")

    # Test
    st.header("🎯 Test Set Performance")
    st.caption("Performance metrics on the held-out test set (best model from hyperparameter search)")
    _render_metrics(test_metrics)

    # Train
    st.header("🏋️‍♂️ Training Set Performance")
    st.caption("Performance metrics on the training set (best model from hyperparameter search)")
    _render_metrics(train_metrics)

    # Complete / All
    st.header("🧮 Overall Performance")
    st.caption("Overall performance metrics combining training and test sets")
    _render_metrics(complete_metrics)


@st.fragment
def render_confusion_matrices(selected_result: TrainingResult):
    """Render confusion matrices for classification tasks.

    Args:
        selected_result: The selected TrainingResult object.

    """
    st.header("🎭 Confusion Matrices")

    # Check if this is a classification task
    if selected_result.run.task not in ["binary", "count_regimes", "density_regimes"]:
        st.info(
            "📊 Confusion matrices are only available for classification tasks "
            "(binary, count_regimes, density_regimes)."
        )
        st.caption("Coming soon for regression tasks: residual plots and error distributions.")
        return

    # Check if confusion matrix data is available
    if selected_result.run.confusion_matrix is None:
        st.warning("⚠️ No confusion matrix data found for this training result.")
        return

    cm = selected_result.run.confusion_matrix

    # Add normalization selection
    st.subheader("Display Options")
    normalize_option = st.radio(
        "Normalization",
        options=["No normalization", "Normalize over True Labels", "Normalize over Predicted Labels"],
        horizontal=True,
        help="Choose how to normalize the confusion matrix values",
    )

    # Map selection to normalization mode
    normalize_map = {
        "No normalization": "none",
        "Normalize over True Labels": "true",
        "Normalize over Predicted Labels": "pred",
    }
    normalize_mode = normalize_map[normalize_option]

    cols = st.columns(3)

    with cols[0]:
        # Test Set Confusion Matrix
        st.subheader("Test Set")
        st.caption("Held-out test set")
        fig_test = plot_confusion_matrix(cm["test"], title="Test Set", normalize=normalize_mode)
        st.plotly_chart(fig_test, width="stretch")
    with cols[1]:
        # Training Set Confusion Matrix
        st.subheader("Training Set")
        st.caption("Training set")
        fig_train = plot_confusion_matrix(cm["train"], title="Training Set", normalize=normalize_mode)
        st.plotly_chart(fig_train, width="stretch")
    with cols[2]:
        # Complete Confusion Matrix
        st.subheader("Complete")
        st.caption("Train + Test sets")
        fig_complete = plot_confusion_matrix(cm["complete"], title="Complete", normalize=normalize_mode)
        st.plotly_chart(fig_complete, width="stretch")


def render_cv_statistics_section(cv_stats: CVMetricStatistics, test_score: float):
    """Render cross-validation statistics for selected metric.

    Args:
        cv_stats: CVMetricStatistics object containing cross-validation statistics.
        test_score: The test set score for the selected metric.

    """
    st.header("📈 Cross-Validation Statistics")
    st.caption("Performance during hyperparameter search (averaged across CV folds)")

    col1, col2, col3, col4, col5 = st.columns(5)

    with col1:
        st.metric("Best Score", f"{cv_stats.best_score:.4f}")
    with col2:
        st.metric("Mean Score", f"{cv_stats.mean_score:.4f}")
    with col3:
        st.metric("Std Dev", f"{cv_stats.std_score:.4f}")
    with col4:
        st.metric("Worst Score", f"{cv_stats.worst_score:.4f}")
    with col5:
        st.metric("Median Score", f"{cv_stats.median_score:.4f}")

    if cv_stats.mean_cv_std is not None:
        st.info(f"**Mean CV Std:** {cv_stats.mean_cv_std:.4f} - Average standard deviation across CV folds")

    # Compare with test metric
    st.subheader("CV vs Test Performance")

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Best CV Score", f"{cv_stats.best_score:.4f}")
    with col2:
        st.metric("Test Score", f"{test_score:.4f}")
    with col3:
        delta = test_score - cv_stats.best_score
        delta_pct = (delta / cv_stats.best_score * 100) if cv_stats.best_score != 0 else 0
        st.metric("Difference", f"{delta:+.4f}", delta=f"{delta_pct:+.2f}%")

    if abs(delta) > cv_stats.std_score:
        st.warning(
            "⚠️ Test performance differs significantly (larger than the CV standard deviation) from CV performance. "
            "This may indicate overfitting or data distribution mismatch between training and test sets."
        )
