"""Data page: Visualization of the data."""

from typing import Literal, cast

import streamlit as st
import xarray as xr
from stopuhr import stopwatch

from entropice.dashboard.sections.alphaearth import render_alphaearth_tab
from entropice.dashboard.sections.arcticdem import render_arcticdem_tab
from entropice.dashboard.sections.areas import render_area_information_tab
from entropice.dashboard.sections.dataset_statistics import render_ensemble_details
from entropice.dashboard.sections.era5 import render_era5_tab
from entropice.dashboard.sections.targets import render_target_information_tab
from entropice.dashboard.utils.loaders import load_training_sets
from entropice.dashboard.utils.stats import DatasetStatistics
from entropice.ml.dataset import DatasetEnsemble
from entropice.utils.types import (
    GridConfig,
    L2SourceDataset,
    TemporalMode,
    grid_configs,
)


def render_dataset_configuration_sidebar() -> DatasetEnsemble:
    """Render dataset configuration selector in sidebar with form.

    Stores the selected ensemble in session state when form is submitted.
    """
    with st.sidebar.form("dataset_config_form"):
        st.header("Dataset Configuration")

        # Grid selection
        grid_options = [gc.display_name for gc in grid_configs]

        grid_level_complete = st.selectbox(
            "Grid Configuration",
            options=grid_options,
            index=0,
            help="Select the grid system and resolution level",
            key="dataset_page_grid",
        )

        # Find the selected grid config
        selected_grid_config: GridConfig = next(gc for gc in grid_configs if gc.display_name == grid_level_complete)

        # Temporal mode selection
        temporal_mode = st.selectbox(
            "Temporal Mode",
            options=cast(list[TemporalMode], ["synopsis", "feature", 2018, 2019, 2020, 2021, 2022, 2023]),
            index=0,
            format_func=lambda x: "Synopsis (mean + trend)"
            if x == "synopsis"
            else "Years-as-Features"
            if x == "feature"
            else f"Year {x}",
            help="Select temporal mode: 'synopsis' for temporal features or specific year",
            key="dataset_page_temporal_mode",
        )

        # Members selection
        st.subheader("Dataset Members")

        all_members = cast(
            list[L2SourceDataset],
            ["AlphaEarth", "ArcticDEM", "ERA5-yearly", "ERA5-seasonal", "ERA5-shoulder"],
        )
        selected_members: list[L2SourceDataset] = []

        for member in all_members:
            if st.checkbox(member, value=True, help=f"Include {member} in the dataset"):
                selected_members.append(member)

        # Form submit button
        submitted = st.form_submit_button(
            "Load Dataset",
            type="primary",
            use_container_width=True,
            disabled=len(selected_members) == 0,
        )

        if not submitted:
            st.info("👆 Click 'Load Dataset' to apply changes.")
            st.stop()

    ensemble = DatasetEnsemble(
        grid=selected_grid_config.grid,
        level=selected_grid_config.level,
        temporal_mode=cast(TemporalMode, temporal_mode),
        members=selected_members,
    )
    return ensemble


def render_dataset_page():
    """Render the Dataset page of the dashboard."""
    st.title("📊 Entropice Data Visualization")
    st.markdown(
        """
        This section allows you to visualize the source datasets and training data
        used in your experiments. Use the sidebar to configure the dataset parameters.
        """
    )

    # Get the selected dataset ensemble from sidebar
    ensemble = render_dataset_configuration_sidebar()

    # Display dataset ID in a styled container
    st.info(f"Loaded Dataset with ID `{ensemble.id()}`")

    st.divider()

    member_datasets = cast(
        dict[L2SourceDataset, xr.Dataset],
        {member: ensemble.read_member(member, lazy=True) for member in ensemble.members},
    )
    # Render dataset statistics section
    stats = DatasetStatistics.from_ensemble(ensemble, member_datasets=member_datasets)
    render_ensemble_details(ensemble.members, stats.members)

    st.divider()

    # Load data and precompute visualizations
    # Preload the grid GeoDataFrame
    grid_gdf = ensemble.read_grid()

    # Create tabs for different data views
    tab_names = ["🎯 Targets", "📐 Areas"]
    # Add tabs for each member based on what's in the ensemble
    if "AlphaEarth" in ensemble.members:
        tab_names.append("🌍 AlphaEarth")
    if "ArcticDEM" in ensemble.members:
        tab_names.append("🏔️ ArcticDEM")
    era5_members = cast(
        list[Literal["ERA5-yearly", "ERA5-seasonal", "ERA5-shoulder"]],
        [m for m in ensemble.members if m.startswith("ERA5")],
    )
    if era5_members:
        tab_names.append("🌡️ ERA5")
    tab_names += ["🔗 Correlations"]
    tabs = st.tabs(tab_names)

    with tabs[0]:
        st.header("🎯 Target Labels Visualization")
        train_data_dict = load_training_sets(ensemble)
        render_target_information_tab(train_data_dict)
    with tabs[1]:
        st.header("📐 Areas Visualization")
        render_area_information_tab(grid_gdf)
    tab_index = 2
    if "AlphaEarth" in ensemble.members:
        with tabs[tab_index]:
            st.header("🌍 AlphaEarth Visualization")
            alphaearth_ds = member_datasets["AlphaEarth"]
            alphaearth_stats = stats.members["AlphaEarth"]
            render_alphaearth_tab(alphaearth_ds, grid_gdf, alphaearth_stats)
        tab_index += 1
    if "ArcticDEM" in ensemble.members:
        with tabs[tab_index]:
            st.header("🏔️ ArcticDEM Visualization")
            arcticdem_ds = member_datasets["ArcticDEM"].compute()
            arcticdem_stats = stats.members["ArcticDEM"]
            render_arcticdem_tab(arcticdem_ds, grid_gdf, arcticdem_stats)
        tab_index += 1
    if era5_members:
        with tabs[tab_index]:
            st.header("🌡️ ERA5 Visualization")
            era5_member_dataset = {m: member_datasets[m] for m in era5_members}
            era5_member_stats = {m: stats.members[m] for m in era5_members}
            render_era5_tab(era5_member_dataset, grid_gdf, era5_member_stats)
        tab_index += 1

    with tabs[tab_index]:
        st.header("🔗 Cross-Dataset Correlation Analysis")
        # Extract grid area series
        #     grid_area_series = grid_gdf.set_index("cell_id")["area_km2"] if "area_km2" in grid_gdf.columns else None
        #    render_correlations_tab(member_datasets, grid_area_series=grid_area_series)
    st.balloons()
    stopwatch.summary()
