"""Section for visualizing experiment inference maps."""

import streamlit as st

from entropice.dashboard.plots.experiment_comparison import create_inference_maps
from entropice.dashboard.utils.loaders import TrainingResult
from entropice.utils.types import GridConfig


@st.fragment
def render_inference_maps_section(
    experiment_name: str,
    training_results: list[TrainingResult],
) -> None:
    """Render the inference maps section.

    Args:
        experiment_name: Name of the experiment
        training_results: List of training results for the experiment

    """
    st.header("🗺️ Inference Maps Analysis")

    st.markdown(
        """
        Visualize the mean and uncertainty (standard deviation) of model predictions across the Arctic region.
        The maps show the spatial distribution of predictions aggregated across multiple training runs.
        """
    )

    if not training_results:
        st.info("No training results available for inference maps.")
        return

    # Extract unique grid configurations from training results
    available_grid_configs = sorted(
        {GridConfig.from_grid_level((tr.run.dataset.grid, tr.run.dataset.level)) for tr in training_results},
        key=lambda gc: gc.sort_key,
    )
    available_tasks = sorted({tr.run.task for tr in training_results})
    available_targets = sorted({tr.run.target for tr in training_results})
    available_models = sorted({tr.run.model_type for tr in training_results})

    # Create form for selecting parameters
    with st.form("inference_map_form"):
        st.subheader("Map Configuration")

        col1, col2 = st.columns(2)

        with col1:
            selected_grid_config = st.selectbox(
                "Grid Configuration",
                options=available_grid_configs,
                format_func=lambda gc: gc.display_name,
                help="Select the grid type and resolution level for the inference map",
            )

            selected_task = st.selectbox(
                "Task",
                options=available_tasks,
                help="Select the prediction task",
            )

        st.subheader("Filters")

        col3, col4 = st.columns(2)

        with col3:
            selected_targets = st.multiselect(
                "Target Datasets",
                options=available_targets,
                default=available_targets,
                help="Filter by target datasets (select all to include all)",
            )

        with col4:
            selected_models = st.multiselect(
                "Model Types",
                options=available_models,
                default=available_models,
                help="Filter by model types (select all to include all)",
            )

        submit_button = st.form_submit_button("Generate Maps", type="primary")

    if submit_button:
        # Extract grid and level from selected config
        selected_grid = selected_grid_config.grid
        selected_level = selected_grid_config.level

        # Filter training results based on selections
        filtered_results = [
            tr
            for tr in training_results
            if tr.run.dataset.grid == selected_grid
            and tr.run.dataset.level == selected_level
            and tr.run.task == selected_task
            and tr.run.target in selected_targets
            and tr.run.model_type in selected_models
        ]

        if not filtered_results:
            st.warning("No training results match the selected criteria. Please adjust your selections and try again.")
            return

        st.success(f"Found {len(filtered_results)} training runs matching the criteria.")

        # Display metadata about selected runs
        with st.expander("View Selected Training Runs"):
            run_info = []
            for tr in filtered_results:
                info = tr.display_info
                run_info.append(
                    {
                        "Task": info.task,
                        "Target": info.target,
                        "Model": info.model,
                        "Grid": f"{info.grid}_{info.level}",
                        "Created": info.timestamp.strftime("%Y-%m-%d %H:%M"),
                    }
                )
            st.dataframe(run_info, width="stretch")

        # Calculate inference maps
        with st.spinner("Calculating inference maps from predictions..."):
            try:
                from entropice.dashboard.utils.loaders import TrainingResult

                inference_gdf = TrainingResult.calculate_inference_maps(filtered_results)

                st.info(
                    f"Generated inference map with {len(inference_gdf):,} grid cells. "
                    + (
                        "Using point-based rendering (>50k cells)."
                        if len(inference_gdf) > 50000
                        else "Using polygon-based rendering."
                    )
                )

            except AssertionError as e:
                st.error(f"Error calculating inference maps: {e}")
                return
            except Exception as e:
                st.error(f"Unexpected error calculating inference maps: {e}")
                st.exception(e)
                return

        # Generate maps
        with st.spinner("Generating cartographic visualizations..."):
            try:
                deck_mean, deck_std = create_inference_maps(
                    inference_gdf,
                    grid=selected_grid,
                    level=selected_level,
                    task=selected_task,
                )

                # Display maps
                st.subheader("📍 Mean Prediction Map")
                st.pydeck_chart(deck_mean, use_container_width=True)

                st.subheader("📍 Uncertainty Map (Standard Deviation)")
                st.pydeck_chart(deck_std, use_container_width=True)

                # Display statistics
                st.subheader("📊 Inference Statistics")
                col1, col2, col3, col4 = st.columns(4)

                with col1:
                    st.metric("Grid Cells", f"{len(inference_gdf):,}")

                with col2:
                    st.metric(
                        "Mean Prediction",
                        f"{inference_gdf['mean_prediction'].mean():.4f}",
                    )

                with col3:
                    st.metric(
                        "Avg Uncertainty",
                        f"{inference_gdf['std_prediction'].mean():.4f}",
                    )

                with col4:
                    st.metric(
                        "Max Uncertainty",
                        f"{inference_gdf['std_prediction'].max():.4f}",
                    )

            except Exception as e:
                st.error(f"Error generating maps: {e}")
                st.exception(e)
                return
