"""ArcticDEM terrain features dashboard section."""

from typing import cast

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

from entropice.dashboard.plots.terrain import (
    create_aspect_rose_diagram,
    create_correlation_heatmap,
    create_slope_aspect_scatter,
    create_terrain_distribution_plot,
    create_terrain_map,
)
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


@st.fragment
def _render_terrain_map(arcticdem_ds: xr.Dataset, grid_gdf: gpd.GeoDataFrame):
    """Visualize spatial distribution of terrain features.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features
        grid_gdf: GeoDataFrame with grid cell geometries

    """
    st.subheader("Spatial Distribution of Terrain Features")

    # Get available features and aggregations
    available_vars = list(arcticdem_ds.data_vars)
    aggs = list(arcticdem_ds.coords["aggregations"].values)

    # Feature selection with grouping
    feature_groups = {
        "TPI (Topographic Position Index)": ["tpi_small", "tpi_medium", "tpi_large"],
        "TRI (Terrain Ruggedness Index)": ["tri_small", "tri_medium", "tri_large"],
        "VRM (Vector Ruggedness Measure)": ["vrm_small", "vrm_medium", "vrm_large"],
        "Slope & Curvature": ["slope", "curvature"],
        "Aspect": ["aspect"],
    }

    # Flatten to get all features
    all_features = [f for features in feature_groups.values() for f in features if f in available_vars]

    cols = st.columns([2, 2, 1])
    with cols[0]:
        # Feature selection with nice formatting
        selected_feature = st.selectbox(
            "Terrain Feature",
            options=all_features,
            format_func=lambda x: x.replace("_", " ").title(),
            key="terrain_feature",
        )

    with cols[1]:
        # Aggregation selection
        selected_agg = st.selectbox(
            "Aggregation",
            options=aggs,
            key="terrain_agg",
        )

    with cols[2]:
        st.write("\n")
        make_3d_map = cast(bool, st.toggle("3D Map", value=True, key="terrain_map_3d"))

    # Extract the data
    if selected_feature not in arcticdem_ds.data_vars:
        st.error(f"Feature {selected_feature} not found in dataset")
        return

    terrain_data = arcticdem_ds[selected_feature].sel(aggregations=selected_agg)
    terrain_series = terrain_data.to_series()

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

    # Create map
    map_deck = create_terrain_map(terrain_series, grid_gdf, selected_feature, make_3d_map)
    st.pydeck_chart(map_deck)

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

        values = terrain_series.dropna()
        if len(values) > 0:
            vmin, vmax = values.min(), values.max()
            vmean = values.mean()
            vstd = values.std()

            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.metric("Min", f"{vmin:.2f}")
            with col2:
                st.metric("Mean", f"{vmean:.2f}")
            with col3:
                st.metric("Max", f"{vmax:.2f}")
            with col4:
                st.metric("Std Dev", f"{vstd:.2f}")

            # Color scale visualization
            if "aspect" in selected_feature.lower():
                cmap = get_cmap("aspect")
            else:
                cmap = get_cmap("terrain")

            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:.2f}</span>
                        <span style="font-size: 0.9em;">{vmax:.2f}</span>
                    </div>
                    <div style="{gradient_style}"></div>
                </div>
                """,
                unsafe_allow_html=True,
            )
        else:
            st.warning("No data available for the selected feature and aggregation")


@st.fragment
def _render_terrain_distributions(arcticdem_ds: xr.Dataset):
    """Display distribution plots for terrain features.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features

    """
    st.subheader("Terrain Feature Distributions")

    # Get available features
    available_vars = list(arcticdem_ds.data_vars)

    # Feature group selection
    feature_groups = {
        "Topographic Position (TPI)": ["tpi_small", "tpi_medium", "tpi_large"],
        "Terrain Ruggedness (TRI)": ["tri_small", "tri_medium", "tri_large"],
        "Vector Ruggedness (VRM)": ["vrm_small", "vrm_medium", "vrm_large"],
        "All Scale-Invariant": ["slope", "aspect", "curvature"],
    }

    selected_group = st.selectbox(
        "Feature Group",
        options=list(feature_groups.keys()),
        key="terrain_dist_group",
    )

    features_to_plot = [f for f in feature_groups[selected_group] if f in available_vars]

    if features_to_plot:
        # Check if subsampling will occur
        n_cells = len(arcticdem_ds.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_terrain_distribution_plot(arcticdem_ds, features_to_plot)
        st.plotly_chart(fig, width="stretch")

        st.markdown(
            f"""
            Distribution of **{selected_group}** features across different aggregation types.
            Violin plots show the full distribution with embedded box plots for quartiles.
            """
        )
    else:
        st.warning(f"No features available for {selected_group}")


@st.fragment
def _render_terrain_correlations(arcticdem_ds: xr.Dataset):
    """Display correlation analysis for terrain features.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features

    """
    st.subheader("Feature Correlation Analysis")

    # Get available features
    available_vars = list(arcticdem_ds.data_vars)
    aggs = list(arcticdem_ds.coords["aggregations"].values)

    # Select aggregation for correlation
    selected_agg = st.selectbox(
        "Aggregation Type",
        options=aggs,
        key="terrain_corr_agg",
    )

    # Select which features to include
    all_features = [str(f) for f in available_vars if f in available_vars]

    if len(all_features) >= 2:
        fig = create_correlation_heatmap(arcticdem_ds, all_features, selected_agg)
        st.plotly_chart(fig, width="stretch")

        st.markdown(
            """
            Correlation heatmap shows relationships between different terrain features.
            Strong positive correlations (red) or negative correlations (blue) can indicate
            related terrain characteristics.
            """
        )
    else:
        st.warning("Need at least 2 features for correlation analysis")


@st.fragment
def _render_slope_aspect_analysis(arcticdem_ds: xr.Dataset, grid_gdf: gpd.GeoDataFrame):
    """Specialized visualization for slope and aspect relationships.

    Args:
        arcticdem_ds: Xarray Dataset containing ArcticDEM terrain features
        grid_gdf: GeoDataFrame with grid cell geometries

    """
    st.subheader("Slope & Aspect Analysis")

    # Check if slope and aspect are available
    if "slope" not in arcticdem_ds.data_vars or "aspect" not in arcticdem_ds.data_vars:
        st.warning("Slope and/or aspect data not available in this dataset")
        return

    # Get aggregations
    aggs = list(arcticdem_ds.coords["aggregations"].values)

    # Select aggregation
    selected_agg = st.selectbox(
        "Aggregation Type",
        options=aggs,
        key="slope_aspect_agg",
    )

    # Extract slope and aspect data
    slope_data = arcticdem_ds["slope"].sel(aggregations=selected_agg).values
    aspect_data = arcticdem_ds["aspect"].sel(aggregations=selected_agg).values

    # Check if subsampling will occur
    n_values = len(slope_data.flatten())
    if n_values > 50000:
        st.info(
            f"📊 **Dataset subsampled:** Using 50,000 randomly selected values out of {n_values:,} "
            "for performance. Distributions remain representative."
        )

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

    with col1:
        st.markdown("**Aspect Rose Diagram**")
        fig_rose = create_aspect_rose_diagram(aspect_data)
        st.plotly_chart(fig_rose, width="stretch")

        st.markdown(
            """
            The rose diagram shows the directional distribution of terrain aspect.
            Each bar represents the frequency of slopes facing that direction.
            """
        )

    with col2:
        st.markdown("**Slope vs Aspect Distribution**")
        fig_scatter = create_slope_aspect_scatter(slope_data, aspect_data)
        st.plotly_chart(fig_scatter, width="stretch")

        st.markdown(
            """
            This 2D histogram shows how slope steepness relates to aspect direction.
            Patterns can reveal preferential slope orientations (e.g., due to prevailing winds or sun exposure).
            """
        )


def render_arcticdem_tab(arcticdem_ds: xr.Dataset, grid_gdf: gpd.GeoDataFrame, arcticdem_stats: MemberStatistics):
    """Render the ArcticDEM visualization tab.

    Args:
        arcticdem_ds: The ArcticDEM dataset member, lazily loaded.
        grid_gdf: GeoDataFrame with grid cell geometries
        arcticdem_stats: Statistics for the ArcticDEM member.

    """
    # Render different visualizations
    with st.expander("ArcticDEM Statistics", expanded=True):
        render_member_details("ArcticDEM", arcticdem_stats)

    st.divider()

    _render_terrain_map(arcticdem_ds, grid_gdf)

    # st.divider()

    # _render_terrain_distributions(arcticdem_ds)

    st.divider()

    _render_terrain_correlations(arcticdem_ds)

    st.divider()

    _render_slope_aspect_analysis(arcticdem_ds, grid_gdf)
