"""Inference page: Visualization of model inference results across the study region."""

import streamlit as st
from stopuhr import stopwatch

from entropice.dashboard.utils.loaders import TrainingResult, load_all_training_results


def render_sidebar_selection(training_results: list[TrainingResult]) -> TrainingResult:
    """Render sidebar for training run selection.

    Args:
        training_results: List of available TrainingResult objects.

    Returns:
        Selected TrainingResult object.

    """
    st.header("Select Training Run")

    # Create selection options with task-first naming
    training_options: dict[str, TrainingResult] = {
        tr.display_info.get_display_name("task_first"): tr for tr in training_results
    }

    selected_name = st.selectbox(
        "Training Run",
        options=list(training_options.keys()),
        index=0,
        help="Select a training run to view inference results",
        key="inference_run_select",
    )

    selected_result = training_options[selected_name]

    st.divider()

    # Show run information in sidebar
    st.subheader("Run Information")

    st.markdown(f"**Task:** {selected_result.run.task.capitalize()}")
    st.markdown(f"**Model:** {selected_result.run.model_type.upper()}")
    st.markdown(f"**Grid:** {selected_result.run.dataset.grid.capitalize()}")
    st.markdown(f"**Level:** {selected_result.run.dataset.level}")
    st.markdown(f"**Target:** {selected_result.run.target.replace('darts_', '')}")
    return selected_result


def render_inference_page():
    """Render the Inference page of the dashboard."""
    st.title("🗺️ Inference Results")

    st.markdown(
        """
        Explore spatial predictions from trained models across the Arctic permafrost region.
        Select a training run from the sidebar to visualize prediction maps, class distributions,
        and spatial coverage statistics.
        """
    )

    # Load all available training results
    training_results = load_all_training_results()

    if not training_results:
        st.warning("No training results found. Please run some training experiments first.")
        st.info("Run training using: `pixi run python -m entropice.ml.training`")
        return

    st.success(f"Found **{len(training_results)}** training result(s)")

    st.divider()

    # Sidebar: Training run selection
    with st.sidebar:
        selected_result = render_sidebar_selection(training_results)

    # Main content area
    # Load predictions
    with st.spinner("Loading inference results..."):
        # Columns: Index(['cell_id', 'predicted', 'geometry'], dtype='object')
        predictions_gdf = selected_result.run.predictions
        task = selected_result.run.task

    # TODO: Implement the sections
    # Map, optionally 3D
    # Some statistics about the predictions
    # Class Distribution for classification tasks
    # Distribution of predicted values for regression tasks

    st.balloons()
    stopwatch.summary()
