"""AlphaEarth embeddings 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.embeddings import (
    create_embedding_distribution_plot,
    create_embedding_map,
    create_embedding_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_band_agg_options(embeddings: xr.Dataset):
    """Get band and aggregation selection options from user."""
    bands = embeddings["band"].values.tolist()
    aggregations = embeddings["agg"].values.tolist()

    cols = st.columns([2, 2])
    with cols[0]:
        band = st.selectbox(
            "Select Embedding Band",
            options=bands,
            index=0,
            help="Select which embedding band to visualize on the map.",
            key="embedding_band_select",
        )
    with cols[1]:
        aggregation = st.selectbox(
            "Select Aggregation Method",
            options=aggregations,
            index=0,
            help="Select the aggregation method for the embeddings to visualize.",
            key="embedding_agg_select",
        )

    return band, aggregation


@st.fragment
def _render_embedding_map(embedding_values: xr.DataArray, grid_gdf: gpd.GeoDataFrame):
    st.subheader("AlphaEarth Embedding Map")

    st.markdown(
        """
        This interactive map visualizes the spatial distribution of the selected embedding band across
        the Arctic region. Each grid cell is colored according to its embedding value, revealing spatial
        patterns in the satellite imagery features. High-resolution embeddings can indicate areas with
        distinctive characteristics that may be relevant for RTS detection.

        **Map controls:**
        - **Hover** over cells to see exact embedding values
        - **3D mode**: Elevation represents embedding magnitude - higher areas have larger values
        - **Rotate** (3D mode): Hold Ctrl/Cmd and drag to rotate the view
        - **Zoom/Pan**: Scroll to zoom, click and drag to pan
        """
    )

    if "year" in embedding_values.dims or "year" in embedding_values.coords:
        year_values = embedding_values["year"].values.tolist()
        year = st.slider(
            "Select Year",
            min_value=int(min(year_values)),
            max_value=int(max(year_values)),
            value=int(max(year_values)),
            step=1,
            help="Select the year for which to visualize the embeddings.",
        )
        embedding_values = embedding_values.sel(year=year)
    make_3d_map = cast(bool, st.toggle("3D Map", value=True, key="embedding_map_3d"))

    # Check if subsampling will occur
    n_cells = len(embedding_values["cell_ids"])
    if n_cells > 50000:
        st.info(f"🗺️ **Map subsampled:** Displaying 50,000 randomly selected cells out of {n_cells:,} for performance.")

    map_deck = create_embedding_map(
        embedding_values=embedding_values,
        grid_gdf=grid_gdf,
        make_3d_map=make_3d_map,
    )

    st.pydeck_chart(map_deck, width="stretch")

    # Add legend
    with st.expander("Legend", expanded=True):
        st.markdown("**Embedding Value**")

        # Get the actual values to show accurate min/max (same as in the map function)
        values = embedding_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}"

        # Get the same colormap used in the map
        cmap = get_cmap("AlphaEarth")
        # Sample 4 colors from the colormap to create the gradient
        gradient_colors = [mcolors.to_hex(cmap(i)) for i in [0.0, 0.33, 0.67, 1.0]]
        gradient_css = ", ".join(gradient_colors)

        # Create a simple gradient legend
        st.markdown(
            f'<div style="display: flex; align-items: center; margin-top: 10px; margin-bottom: 10px;">'
            f'<span style="margin-right: 10px;">{vmin_str}</span>'
            f'<div style="flex: 1; height: 20px; background: linear-gradient(to right, '
            f'{gradient_css}); border: 1px solid #ccc;"></div>'
            f'<span style="margin-left: 10px;">{vmax_str}</span>'
            f"</div>",
            unsafe_allow_html=True,
        )

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

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


def _render_trend(embeddin_values: xr.DataArray):
    st.subheader("AlphaEarth Embedding Trends Over Time")

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

        **Understanding the plot:**
        - **Mean line** (dark green): Average embedding value across all grid cells for each year
        - **10th-90th percentile band** (light green): Range containing 80% of the values, showing
          typical variation
        - **Min/Max range** (gray): Full extent of values, highlighting outliers
        """
    )

    # Show dataset filtering info
    band_val = embeddin_values["band"].values.item() if embeddin_values["band"].size == 1 else "multiple"
    agg_val = embeddin_values["agg"].values.item() if embeddin_values["agg"].size == 1 else "multiple"
    st.caption(
        f"📊 **Dataset selection:** Band `{band_val}`, Aggregation `{agg_val}` "
        f"({len(embeddin_values['year'])} years, {len(embeddin_values['cell_ids']):,} cells)"
    )

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

    fig = create_embedding_trend_plot(embedding_values=embeddin_values)
    st.plotly_chart(fig, width="stretch")


def _render_distribution(embeddin_values: xr.DataArray):
    st.subheader("AlphaEarth Embedding Distribution")

    st.markdown(
        """
        This plot shows the statistical distribution of embedding values across all 64 embedding
        dimensions (bands). AlphaEarth embeddings are learned representations from satellite imagery,
        with each band capturing different aspects of the landscape (e.g., vegetation, terrain, ice
        cover, land use).

        **Understanding the plot:**
        - **X-axis**: Embedding bands (A00-A63), each representing a learned feature from satellite
          imagery
        - **Mean line** (dark green): Average value across all grid cells for each band
        - **10th-90th percentile band** (light green): Central distribution of values, excluding
          outliers
        - **Min/Max range** (gray): Full value range showing extreme values

        Different bands may capture different landscape features - bands with higher variance often
        represent more spatially heterogeneous characteristics.
        """
    )

    # Show dataset filtering info
    agg_val = embeddin_values["agg"].values.item() if embeddin_values["agg"].size == 1 else "multiple"
    n_bands = len(embeddin_values["band"])
    n_cells = len(embeddin_values["cell_ids"])
    st.caption(f"📊 **Dataset selection:** Aggregation `{agg_val}` ({n_bands} bands, {n_cells:,} cells)")

    # 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."
        )

    fig = create_embedding_distribution_plot(embedding_values=embeddin_values)
    st.plotly_chart(fig, width="stretch")


@st.fragment
def render_alphaearth_tab(embeddings: xr.Dataset, grid_gdf: gpd.GeoDataFrame, member_stats: MemberStatistics):
    """Render the AlphaEarth visualization tab.

    Args:
        embeddings: The AlphaEarth dataset member, lazily loaded.
        grid_gdf: GeoDataFrame with grid cell geometries
        member_stats: Statistics for the AlphaEarth member.

    """
    # Render different visualizations
    with st.expander("AlphaEarth Embedding Statistics", expanded=True):
        render_member_details("AlphaEarth", member_stats)

    st.divider()

    band, aggregation = _get_band_agg_options(embeddings)
    embedding_values = embeddings["embeddings"].sel(agg=aggregation).compute()

    _render_distribution(embedding_values)
    st.divider()

    if "year" in embedding_values.dims or "year" in embedding_values.coords:
        _render_trend(embedding_values.sel(band=band))
        st.divider()

    _render_embedding_map(embedding_values.sel(band=band), grid_gdf)
