"""Cross-dataset correlation analysis dashboard section."""

from typing import Any, cast

import pandas as pd
import streamlit as st
import xarray as xr

from entropice.dashboard.plots.correlations import (
    create_dendrogram_plot,
    create_feature_variance_plot,
    create_full_correlation_heatmap,
    create_matplotlib_correlation_heatmap,
    create_mutual_information_matrix,
    create_pca_biplot,
    select_top_variable_features,
)
from entropice.utils.types import L2SourceDataset


def _get_aggregation_dimensions(ds: xr.Dataset) -> list[str]:
    """Get aggregation dimension names from a dataset.

    Args:
        ds: Xarray Dataset

    Returns:
        List of aggregation dimension names

    """
    return [str(dim) for dim in ds.dims if dim in ("agg", "aggregations")]


def _set_all_aggregations_corr(
    member_datasets: dict[L2SourceDataset, xr.Dataset],
    members_with_aggs: list[L2SourceDataset],
    member_agg_dims: dict[L2SourceDataset, list[str]],
    selected: bool,
):
    """Set all aggregation checkboxes to selected or deselected state.

    Args:
        member_datasets: Dictionary mapping members to their loaded datasets
        members_with_aggs: List of members that have aggregations
        member_agg_dims: Dictionary mapping members to aggregation dimensions
        selected: True to select all, False to deselect all

    """
    for member in members_with_aggs:
        ds = member_datasets.get(member)
        if ds is None:
            continue
        agg_dims = member_agg_dims.get(member, [])
        for agg_dim in agg_dims:
            if agg_dim in ds.dims:
                agg_values = ds.coords[agg_dim].to_numpy().tolist()
                for val in agg_values:
                    st.session_state[f"corr_agg_{member}_{agg_dim}_{val}"] = selected


def _set_median_only_aggregations_corr(
    member_datasets: dict[L2SourceDataset, xr.Dataset],
    members_with_aggs: list[L2SourceDataset],
    member_agg_dims: dict[L2SourceDataset, list[str]],
):
    """Set only median aggregations to selected, deselect all others.

    Args:
        member_datasets: Dictionary mapping members to their loaded datasets
        members_with_aggs: List of members that have aggregations
        member_agg_dims: Dictionary mapping members to aggregation dimensions

    """
    for member in members_with_aggs:
        ds = member_datasets.get(member)
        if ds is None:
            continue
        agg_dims = member_agg_dims.get(member, [])
        for agg_dim in agg_dims:
            if agg_dim in ds.dims:
                agg_values = ds.coords[agg_dim].to_numpy().tolist()
                for val in agg_values:
                    # Select only if value is or contains 'median'
                    is_median = str(val).lower() == "median" or "median" in str(val).lower()
                    st.session_state[f"corr_agg_{member}_{agg_dim}_{val}"] = is_median


def _render_member_aggregation_checkboxes(
    member: L2SourceDataset,
    ds: xr.Dataset,
    agg_dims: list[str],
    column: Any,
    dimension_filters: dict[str, dict[str, list[str]]],
) -> dict[str, dict[str, list[str]]]:
    """Render checkboxes for a single member's aggregations."""
    with column:
        st.markdown(f"**{member}:**")

        for agg_dim in agg_dims:
            if agg_dim not in ds.dims:
                continue

            # Get coordinate values
            agg_values = [str(v) for v in ds.coords[agg_dim].to_numpy().tolist()]

            st.markdown(f"*{agg_dim}:*")
            selected_vals = []

            for val in agg_values:
                key = f"corr_agg_{member}_{agg_dim}_{val}"

                if st.checkbox(val, value=True, key=key, help=f"Include {val} from {agg_dim}"):
                    selected_vals.append(val)

            # Store selected values if not all selected
            if selected_vals and len(selected_vals) < len(agg_values):
                if member not in dimension_filters:
                    dimension_filters[member] = {}
                dimension_filters[member][agg_dim] = selected_vals

    return dimension_filters


def _render_aggregation_selection(
    member_datasets: dict[L2SourceDataset, xr.Dataset],
) -> dict[str, dict[str, list[str]]]:
    """Render aggregation selection controls for members that have aggregations.

    Args:
        member_datasets: Dictionary mapping member names to their xarray Datasets

    Returns:
        Dictionary mapping member names to dimension filters

    """
    # Find members with aggregations
    members_with_aggs = []
    member_agg_dims = {}

    for member, ds in member_datasets.items():
        agg_dims = _get_aggregation_dimensions(ds)
        if agg_dims:
            members_with_aggs.append(member)
            member_agg_dims[member] = agg_dims

    if not members_with_aggs:
        return {}

    st.markdown("#### Select Aggregations")
    st.markdown(
        "Select which spatial aggregations to include. "
        "This allows you to analyze correlations for specific aggregation types (e.g., mean, median, std)."
    )

    # Add buttons for common selections (outside form to allow state manipulation)
    col_btn1, col_btn2, col_btn3, _ = st.columns([1, 1, 1, 3])

    with col_btn1:
        if st.button("✅ Select All", use_container_width=True, key="corr_agg_select_all"):
            _set_all_aggregations_corr(member_datasets, members_with_aggs, member_agg_dims, selected=True)
    with col_btn2:
        if st.button("📊 Median Only", use_container_width=True, key="corr_agg_median_only"):
            _set_median_only_aggregations_corr(member_datasets, members_with_aggs, member_agg_dims)
    with col_btn3:
        if st.button("❌ Deselect All", use_container_width=True, key="corr_agg_deselect_all"):
            _set_all_aggregations_corr(member_datasets, members_with_aggs, member_agg_dims, selected=False)

    # Render form with checkboxes
    with st.form("correlation_aggregation_selection_form"):
        dimension_filters: dict[str, dict[str, list[str]]] = {}

        # Create columns for each member
        member_cols = st.columns(len(members_with_aggs))

        for col_idx, member in enumerate(members_with_aggs):
            dimension_filters = _render_member_aggregation_checkboxes(
                member,
                member_datasets[member],
                member_agg_dims[member],
                member_cols[col_idx],
                dimension_filters,
            )

        # Submit button for the form
        submitted = st.form_submit_button("Apply Aggregation Filters", type="primary", use_container_width=True)

        if not submitted:
            st.info("👆 Click 'Apply Aggregation Filters' to update the configuration")
            st.stop()

    return dimension_filters


def _flatten_dataset_to_series(
    ds: xr.Dataset,
    prefix: str = "",
    dimension_filter: dict[str, list[str]] | None = None,
) -> dict[str, pd.Series]:
    """Flatten an xarray Dataset into a dict of pandas Series.

    Handles multi-dimensional variables by creating separate series for each combination
    of non-cell_ids dimensions.

    Args:
        ds: Xarray Dataset to flatten
        prefix: Prefix to add to variable names
        dimension_filter: Optional dict mapping dimension names to lists of allowed values.
            Only combinations matching these values will be included.

    Returns:
        Dictionary mapping variable names to pandas Series with cell_ids as index

    """
    series_dict = {}
    dimension_filter = dimension_filter or {}

    for var_name in ds.data_vars:
        var_data = ds[var_name]

        # Get dimensions other than cell_ids
        other_dims = [dim for dim in var_data.dims if dim != "cell_ids"]

        if len(other_dims) == 0:
            # Simple 1D variable
            series = var_data.to_series()
            full_name = f"{prefix}{var_name}" if prefix else var_name
            series_dict[full_name] = series
        else:
            # Multi-dimensional variable - create series for each combination
            for coord_values in var_data.stack(stacked=other_dims).coords["stacked"].to_numpy():  # noqa: PD013
                # Create selector dict
                if isinstance(coord_values, tuple):
                    selector = dict(zip(other_dims, coord_values))
                else:
                    selector = {other_dims[0]: coord_values}

                # Check if this combination passes the dimension filter
                if dimension_filter:
                    skip = False
                    for dim, val in selector.items():
                        if dim in dimension_filter:
                            # Check if value is in allowed list (convert to string for comparison)
                            if str(val) not in dimension_filter[dim]:
                                skip = True
                                break
                    if skip:
                        continue

                # Extract 1D series
                series = var_data.sel(selector).to_series()

                # Create descriptive name
                suffix = "_".join(str(v) for v in (coord_values if isinstance(coord_values, tuple) else [coord_values]))
                full_name = f"{prefix}{var_name}_{suffix}" if prefix else f"{var_name}_{suffix}"
                series_dict[full_name] = series

    return series_dict


@st.fragment
def _render_correlation_matrix(data_dict: dict[str, pd.Series]):
    """Render correlation matrix visualization.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series

    """
    st.subheader("Correlation Matrix")

    n_features = len(data_dict)

    # Add feature reduction options for large datasets
    reduced_data = data_dict
    if n_features > 500:
        st.warning(
            f"⚠️ **Large dataset detected:** {n_features} features may be too many to visualize effectively. "
            "Consider reducing the feature set using the options below."
        )

        with st.expander("🎯 Feature Reduction Options", expanded=True):
            reduction_method = st.radio(
                "Reduction Strategy",
                options=["top_variable", "none"],
                format_func=lambda x: (
                    "Select Most Variable Features" if x == "top_variable" else "Use All Features (may be slow)"
                ),
                index=0,
                key="corr_reduction_method",
            )

            if reduction_method == "top_variable":
                col1, col2 = st.columns(2)
                with col1:
                    n_top = st.slider(
                        "Number of Features to Keep",
                        min_value=50,
                        max_value=min(1000, n_features),
                        value=min(500, n_features),
                        step=50,
                        key="corr_n_top",
                    )
                with col2:
                    variability_metric = st.selectbox(
                        "Variability Metric",
                        options=["variance", "iqr", "cv"],
                        format_func=lambda x: {
                            "variance": "Variance",
                            "iqr": "Interquartile Range",
                            "cv": "Coefficient of Variation",
                        }[x],
                        key="corr_variability_metric",
                    )

                reduced_data = select_top_variable_features(data_dict, n_features=n_top, method=variability_metric)
                st.success(f"✅ Reduced from {n_features} to {len(reduced_data)} most variable features")

    n_reduced = len(reduced_data)

    # Configuration options
    col1, col2 = st.columns([3, 1])

    with col1:
        method = st.selectbox(
            "Correlation Method",
            options=["pearson", "spearman", "kendall"],
            index=0,
            format_func=lambda x: x.title(),
            help="Pearson: linear relationships | Spearman: monotonic relationships | Kendall: robust to outliers",
            key="corr_method",
        )

    with col2:
        cluster = cast(
            bool,
            st.toggle(
                "Cluster Variables",
                value=n_reduced < 100,
                help="Reorder by hierarchical clustering",
                key="corr_cluster",
            ),
        )

    # Check if subsampling will occur
    df_test = pd.DataFrame(reduced_data).dropna()
    if len(df_test) > 50000:
        st.info(
            f"📊 **Dataset subsampled:** Using 50,000 randomly selected cells out of {len(df_test):,} "
            "for performance. Correlations remain representative."
        )

    # Use matplotlib for large feature sets (> 200), plotly for smaller
    use_matplotlib = n_reduced > 200

    if use_matplotlib:
        st.info(
            f"📈 **Using Matplotlib rendering** for {n_reduced} features "
            "(Plotly would exceed Streamlit's message size limit). "
            "This produces a static image but handles large datasets efficiently."
        )
        fig = create_matplotlib_correlation_heatmap(reduced_data, method=method, cluster=cluster)
        st.pyplot(fig, use_container_width=True)
        # Close figure to free memory
        import matplotlib.pyplot as plt_cleanup

        plt_cleanup.close(fig)
    else:
        fig = create_full_correlation_heatmap(reduced_data, method=method, cluster=cluster)
        st.plotly_chart(fig, width="stretch")

    st.markdown(
        f"""
        **{method.title()} correlation** measures the {"linear" if method == "pearson" else "monotonic"}
        relationship between variables. Values range from -1 (perfect negative correlation) to
        +1 (perfect positive correlation), with 0 indicating no correlation.
        """
    )


@st.fragment
def _render_pca_analysis(data_dict: dict[str, pd.Series]):
    """Render PCA biplot visualization.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series

    """
    st.subheader("Principal Component Analysis (PCA)")

    n_features = len(data_dict)

    # Reduce features if too many for visualization
    reduced_data = data_dict
    if n_features > 500:
        st.warning(
            f"⚠️ **Large dataset:** {n_features} features detected. "
            "Automatically selecting the 500 most variable features for PCA visualization."
        )
        reduced_data = select_top_variable_features(data_dict, n_features=500, method="variance")

    show_loadings = cast(
        bool,
        st.toggle(
            "Show Loading Vectors",
            value=True,
            help="Display how each variable contributes to the principal components",
            key="pca_loadings",
        ),
    )

    fig = create_pca_biplot(reduced_data, n_components=2, show_loadings=show_loadings)
    st.plotly_chart(fig, width="stretch")

    st.markdown(
        """
        **PCA** reduces the dimensionality of the data while preserving variance.
        - **Blue points**: Individual grid cells projected onto the first two principal components
        - **Red arrows**: Loading vectors showing how each variable contributes to the PCs
        - Variables pointing in similar directions are positively correlated
        - Longer arrows indicate variables with higher variance
        """
    )


@st.fragment
def _render_hierarchical_clustering(data_dict: dict[str, pd.Series]):
    """Render hierarchical clustering dendrogram.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series

    """
    st.subheader("Hierarchical Clustering of Variables")

    n_features = len(data_dict)

    # Reduce features if too many
    reduced_data = data_dict
    if n_features > 300:
        st.warning(
            f"⚠️ **Large dataset:** {n_features} features detected. "
            "Automatically selecting the 300 most variable features for dendrogram visualization."
        )
        reduced_data = select_top_variable_features(data_dict, n_features=300, method="variance")

    method = st.selectbox(
        "Distance Metric (based on correlation)",
        options=["pearson", "spearman"],
        index=0,
        format_func=lambda x: x.title(),
        key="dendro_method",
    )

    fig = create_dendrogram_plot(reduced_data, method=method)
    st.plotly_chart(fig, width="stretch")

    st.markdown(
        """
        **Dendrogram** shows hierarchical relationships between variables based on correlation distance.
        - Variables that merge at lower heights are more similar (highly correlated)
        - Distinct clusters suggest groups of related features
        - Useful for identifying redundant features or feature groups
        """
    )


@st.fragment
def _render_mutual_information(data_dict: dict[str, pd.Series]):
    """Render mutual information matrix.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series

    """
    st.subheader("Mutual Information Analysis")

    n_features = len(data_dict)

    # Mutual information is very computationally expensive - hard limit
    if n_features > 200:
        st.error(
            f"❌ **Too many features:** Mutual information analysis with {n_features} features "
            "would be too computationally expensive. Please reduce to ≤200 features using the "
            "variable selection or aggregation filters."
        )
        return

    st.warning(
        "⚠️ **Computationally intensive**: This analysis may take some time. "
        "Dataset is subsampled to 10,000 cells for performance."
    )

    fig = create_mutual_information_matrix(data_dict)
    st.plotly_chart(fig, width="stretch")

    st.markdown(
        """
        **Mutual Information** captures both linear and non-linear relationships between variables.
        - Unlike correlation, MI can detect complex dependencies
        - Higher values indicate stronger information sharing
        - Particularly useful for identifying non-linear relationships that correlation might miss
        """
    )


@st.fragment
def _render_variance_analysis(data_dict: dict[str, pd.Series]):
    """Render feature variance and spread analysis.

    Args:
        data_dict: Dictionary mapping variable names to pandas Series

    """
    st.subheader("Feature Variance and Spread")

    fig = create_feature_variance_plot(data_dict)
    st.plotly_chart(fig, width="stretch")

    st.markdown(
        """
        **Variance metrics** show the spread and variability of each feature:
        - **Standard Deviation**: Average distance from the mean
        - **IQR** (Interquartile Range): Spread of the middle 50% of data (robust to outliers)
        - **Range**: Difference between max and min values
        - **CV** (Coefficient of Variation): Normalized variability (std/mean)

        Features with very low variance may provide little information for modeling.
        """
    )


def _get_variable_groups(all_series: dict[str, pd.Series]) -> dict[str, list[str]]:
    """Group variables by data source.

    Args:
        all_series: Dictionary of all variable series

    Returns:
        Dictionary mapping source names to lists of variable names

    """
    var_groups = {}
    for var_name in all_series.keys():
        source = var_name.split("_")[0]
        if source not in var_groups:
            var_groups[source] = []
        var_groups[source].append(var_name)
    return var_groups


def _render_variable_selection(all_series: dict[str, pd.Series]) -> dict[str, pd.Series]:
    """Render variable selection UI and return filtered data.

    Args:
        all_series: Dictionary of all variable series

    Returns:
        Filtered dictionary with only selected variables

    """
    var_groups = _get_variable_groups(all_series)
    selected_vars = []

    # Create checkboxes for each group
    for source, vars_in_group in sorted(var_groups.items()):
        with st.container():
            col1, col2 = st.columns([1, 4])
            with col1:
                select_all = st.checkbox(f"**{source}** ({len(vars_in_group)})", value=True, key=f"select_{source}")

            with col2:
                if select_all:
                    selected_vars.extend(vars_in_group)
                else:
                    # Show first few variables as examples
                    st.caption(", ".join(vars_in_group[:3]) + ("..." if len(vars_in_group) > 3 else ""))

    # Filter data dict
    filtered_data = {k: v for k, v in all_series.items() if k in selected_vars}

    if len(filtered_data) < 2:
        st.warning("⚠️ Please select at least 2 variables for correlation analysis")
        st.stop()

    st.info(f"**Selected {len(filtered_data)} variables** for analysis")
    return filtered_data


@st.fragment
def render_correlations_tab(
    member_datasets: dict[L2SourceDataset, xr.Dataset],
    grid_area_series: pd.Series | None = None,
):
    """Render the cross-dataset correlation analysis tab.

    Args:
        member_datasets: Dictionary mapping member names to their xarray Datasets
        grid_area_series: Optional pandas Series with cell_ids as index and area values

    """
    st.markdown(
        """
        This section provides comprehensive analysis of relationships and similarities
        between all variables across different data sources (ArcticDEM, ERA5, AlphaEarth, etc.).
        """
    )

    # Flatten all datasets into a single dict of pandas Series
    with st.spinner("Preparing data for correlation analysis..."):
        # First, compute all datasets if needed
        computed_datasets = {}
        for member_name, ds in member_datasets.items():
            # Compute if lazy
            if any(isinstance(v.data, type(ds)) for v in ds.data_vars.values()):
                with st.spinner(f"Loading {member_name} data..."):
                    ds = ds.compute()
            computed_datasets[member_name] = ds

    # Render aggregation selection
    with st.expander("🔧 Aggregation Selection", expanded=False):
        dimension_filters = _render_aggregation_selection(computed_datasets)

        if dimension_filters:
            st.info(
                f"**Filtering applied:** {sum(len(dims) for dims in dimension_filters.values())} "
                f"dimension filter(s) across {len(dimension_filters)} data source(s)"
            )

    # Now flatten with filters applied
    with st.spinner("Flattening datasets..."):
        all_series = {}

        # Add grid area if provided
        if grid_area_series is not None:
            all_series["Grid_Cell_Area_km2"] = grid_area_series

        # Process each member dataset with its dimension filter
        for member_name, ds in computed_datasets.items():
            prefix = f"{member_name}_"
            member_filter = dimension_filters.get(member_name, None)
            member_series = _flatten_dataset_to_series(ds, prefix=prefix, dimension_filter=member_filter)
            all_series.update(member_series)

    st.success(f"✅ Loaded {len(all_series)} variables from {len(member_datasets)} data sources")

    # Variable selection
    with st.expander("🔧 Variable Selection", expanded=False):
        st.markdown("Select which variables to include in the correlation analysis.")
        filtered_data = _render_variable_selection(all_series)

    # Create tabs for different analyses
    analysis_tabs = st.tabs(
        [
            "📊 Correlation Matrix",
            "🔬 PCA Biplot",
            "🌳 Hierarchical Clustering",
            "📈 Variance Analysis",
            "🔗 Mutual Information",
        ]
    )

    with analysis_tabs[0]:
        _render_correlation_matrix(filtered_data)

    with analysis_tabs[1]:
        _render_pca_analysis(filtered_data)

    with analysis_tabs[2]:
        _render_hierarchical_clustering(filtered_data)

    with analysis_tabs[3]:
        _render_variance_analysis(filtered_data)

    with analysis_tabs[4]:
        _render_mutual_information(filtered_data)
