"""Sections for the inference page."""

import geopandas as gpd
import streamlit as st

from entropice.dashboard.plots.inference import (
    create_coverage_statistics_table,
    create_inference_map,
    create_prediction_distribution_plot,
    create_prediction_statistics_plot,
)
from entropice.utils.types import Task


def render_inference_map_section(predictions_gdf: gpd.GeoDataFrame, task: Task) -> None:
    """Render the inference map section.

    Args:
        predictions_gdf: GeoDataFrame with predictions and geometries
        task: Task type

    """
    st.subheader("📍 Spatial Prediction Map")

    st.markdown(
        """
        Interactive map showing model predictions across the Arctic region.
        Toggle 3D mode to see predictions as elevated terrain.
        """
    )

    # 3D toggle
    col1, col2 = st.columns([1, 4])
    with col1:
        make_3d = st.checkbox("3D View", value=False, key="inference_map_3d")

    with col2:
        if make_3d:
            st.info("🎨 Height represents prediction magnitude. Rotate and tilt the map to explore!")

    # Create and display map
    with st.spinner("Generating map..."):
        deck = create_inference_map(predictions_gdf, task, make_3d)
        st.pydeck_chart(deck)


def render_prediction_statistics_section(
    predictions_gdf: gpd.GeoDataFrame, task: Task, grid_gdf: gpd.GeoDataFrame | None = None
) -> None:
    """Render the prediction statistics section.

    Args:
        predictions_gdf: GeoDataFrame with predictions
        task: Task type
        grid_gdf: Optional GeoDataFrame with cell_id and land_area for RTS area calculations

    """
    st.subheader("📊 Prediction Statistics")

    # Create two columns for statistics
    col1, col2 = st.columns(2)

    with col1:
        st.markdown("#### Coverage Statistics")
        # Create statistics table
        stats_df = create_coverage_statistics_table(predictions_gdf, task, grid_gdf)
        st.dataframe(stats_df, hide_index=True, width="stretch")

    with col2:
        st.markdown("#### Statistical Summary")
        # Create statistics plot (box plot for regression, proportions for classification)
        fig = create_prediction_statistics_plot(predictions_gdf, task)
        st.plotly_chart(fig, width="stretch")


def render_prediction_distribution_section(predictions_gdf: gpd.GeoDataFrame, task: Task) -> None:
    """Render the prediction distribution section.

    Args:
        predictions_gdf: GeoDataFrame with predictions
        task: Task type

    """
    classification_tasks: list[Task] = ["binary", "count_regimes", "density_regimes"]
    is_classification = task in classification_tasks

    if is_classification:
        st.subheader("📈 Class Distribution")
        st.markdown(
            """
            Distribution of predicted classes across all grid cells.
            Shows the count and proportion of each predicted class.
            """
        )
    else:
        st.subheader("📈 Value Distribution")
        st.markdown(
            """
            Distribution of predicted values across all grid cells.
            Histogram shows the frequency of predictions in different value ranges.
            """
        )

    # Create and display distribution plot
    fig = create_prediction_distribution_plot(predictions_gdf, task)
    st.plotly_chart(fig, width="stretch")
