"""ERA5 climate data dashboard section."""

from typing import cast

import geopandas as gpd
import matplotlib.colors as mcolors
import numpy as np
import streamlit as st
import xarray as xr

from entropice.dashboard.plots.climate import (
    create_climate_map,
    create_climate_trend_plot,
)
from entropice.dashboard.sections.dataset_statistics import render_member_details
from entropice.dashboard.utils.colors import get_cmap
from entropice.dashboard.utils.stats import MemberStatistics


def _get_climate_variable_aggregation_season_option(era5_ds: xr.Dataset):
    available_vars = list(era5_ds.data_vars)
    selected_var = st.selectbox("Select Climate Variable", available_vars, key="climate_var_map")
    if "month" in era5_ds.dims:
        months = era5_ds.month.to_numpy()
        selected_month = cast(str, st.selectbox("Select Season", options=months.tolist(), key="climate_month"))
    else:
        selected_month = None
    if "aggregations" in era5_ds.dims:
        aggs = era5_ds.aggregations.to_numpy()
        selected_agg = cast(str, st.selectbox("Select Aggregation Method", options=aggs.tolist(), key="climate_agg"))
    else:
        selected_agg = None
    return selected_var, selected_agg, selected_month


@st.fragment
def _render_climate_variable_map(climate_values: xr.DataArray, grid_gdf: gpd.GeoDataFrame, selected_var: str):
    """Visualize spatial distribution of climate variables.

    Args:
        climate_values: Xarray DataArray containing ERA5 climate features
        grid_gdf: GeoDataFrame with grid cell geometries
        selected_var: Name of the selected climate variable

    """
    st.subheader("Spatial Distribution of Climate Variables")

    if "year" in climate_values.dims:
        years = climate_values.year.to_numpy()
        selected_year = cast(
            int, st.select_slider("Select Year", options=years.tolist(), value=int(years.max()), key="climate_year")
        )
        climate_values = climate_values.sel(year=selected_year)

    # 3D toggle
    make_3d = cast(bool, st.toggle("3D Map", value=True, key="climate_map_3d"))

    # Create map
    n_cells = len(climate_values)
    if n_cells > 50000:
        st.info(f"Showing 50,000 / {n_cells:,} cells for performance")

    deck = create_climate_map(climate_values, grid_gdf, selected_var, make_3d)
    st.pydeck_chart(deck, use_container_width=True)

    # Add legend
    with st.expander("Legend", expanded=True):
        st.markdown(f"**{selected_var.replace('_', ' ').title()}**")

        # Get the actual values to show accurate min/max (same as in the map function)
        values = climate_values.values.flatten()
        values = values[~np.isnan(values)]
        vmin, vmax = np.nanpercentile(values, [2, 98])

        vmin_str = f"{vmin:.4f}"
        vmax_str = f"{vmax:.4f}"

        # Color scale visualization - use appropriate colormap
        cmap = get_cmap(selected_var)

        gradient_colors = [mcolors.to_hex(cmap(i)) for i in [0.0, 0.33, 0.67, 1.0]]
        gradient_css = ", ".join(gradient_colors)

        gradient_style = f"height: 20px; background: linear-gradient(to right, {gradient_css}); border-radius: 4px;"
        st.markdown(
            f"""
            <div style="margin-top: 10px;">
                <div style="display: flex; justify-content: space-between; margin-bottom: 4px;">
                    <span style="font-size: 0.9em;">{vmin_str}</span>
                    <span style="font-size: 0.9em;">{vmax_str}</span>
                </div>
                <div style="{gradient_style}"></div>
            </div>
            """,
            unsafe_allow_html=True,
        )

        st.caption(
            "Color intensity represents climate values from low to high. "
            "Values are normalized using the 2nd-98th percentile range to avoid outliers."
        )

        if make_3d:
            st.markdown("---")
            st.markdown("**3D Elevation:**")
            st.caption("Height represents normalized climate values. Rotate the map by holding Ctrl/Cmd and dragging.")


def _render_climate_temporal_trends(
    climate_values: xr.DataArray, selected_var: str, selected_agg: str | None, selected_month: str | None
):
    """Display temporal trends of climate variables.

    Args:
        climate_values: Xarray DataArray containing ERA5 climate features
        selected_var: Name of the selected climate variable
        selected_agg: Selected aggregation method
        selected_month: Selected month/season

    """
    st.subheader("Climate Variable Temporal Trends")

    if "year" not in climate_values.dims:
        st.info("Temporal trends require yearly dimension")
        return

    st.markdown(
        """
        This visualization shows how climate variables have changed over time across the study area.
        The plot aggregates spatial statistics (mean, percentiles) for each year, revealing temporal
        patterns in climate variables that may correlate with environmental changes.

        **Understanding the plot:**
        - **Mean line** (blue): Average value across all grid cells for each year
        - **10th-90th percentile band** (light blue): Range containing 80% of the values, showing
          typical variation
        - **Min/Max range** (gray): Full extent of values, highlighting outliers
        """
    )
    # Show dataset filtering info
    n_years = len(climate_values["year"])
    n_cells = len(climate_values["cell_ids"])

    agg_info = f", Aggregation: `{selected_agg}`" if selected_agg else ""
    month_info = f", Season: `{selected_month}`" if selected_month else ""
    st.caption(
        f"📊 **Dataset selection:** Variable: `{selected_var}` "
        f"({n_years} years, {n_cells:,} cells{agg_info}{month_info})"
    )

    # Check if subsampling will occur
    if n_cells > 10000:
        st.info(
            f"📊 **Dataset subsampled:** Using 10,000 randomly selected cells out of {n_cells:,} "
            "for performance. Statistics remain representative."
        )

    # Create trend plot
    fig = create_climate_trend_plot(climate_values, selected_var)
    st.plotly_chart(fig, use_container_width=True)


@st.fragment
def render_era5_tab(
    era5_member_datasets: dict[str, xr.Dataset],
    grid_gdf: gpd.GeoDataFrame,
    era5_member_stats: dict[str, MemberStatistics],
):
    """Render the ERA5 visualization tab.

    Args:
        era5_member_datasets: Dictionary of ERA5 member datasets (yearly, seasonal, shoulder)
        grid_gdf: GeoDataFrame with grid cell geometries
        era5_member_stats: Dictionary of MemberStatistics for each ERA5 member.

    """
    if not era5_member_datasets:
        st.warning("No ERA5 data available")
        return

    # Member selection
    available_members = list(era5_member_datasets.keys())

    if len(available_members) == 1:
        selected_member = available_members[0]
        st.info(f"Showing ERA5 data: {selected_member}")
    else:
        selected_member = st.selectbox(
            "Select ERA5 Aggregation Type",
            available_members,
            format_func=lambda x: x.replace("_", " ").title(),
            key="era5_member",
        )

    era5_stats = era5_member_stats[selected_member]
    # Load selected dataset
    era5_ds = era5_member_datasets[selected_member]

    # Render different visualizations
    with st.expander(f"{selected_member.replace('_', ' ').title()} Statistics", expanded=True):
        render_member_details(selected_member, era5_stats)

    st.divider()

    selected_var, selected_agg, selected_month = _get_climate_variable_aggregation_season_option(era5_ds)

    climate_values = era5_ds[selected_var]
    if selected_agg:
        climate_values = climate_values.sel(aggregations=selected_agg)
    if selected_month:
        climate_values = climate_values.sel(month=selected_month)
    climate_values = climate_values.compute()

    if "year" in climate_values.dims:
        st.divider()

        _render_climate_temporal_trends(climate_values, selected_var, selected_agg, selected_month)

    _render_climate_variable_map(climate_values, grid_gdf, selected_var)
