"""Feature importance analysis section for experiment comparison."""

import pandas as pd
import streamlit as st

from entropice.dashboard.plots.experiment_comparison import (
    create_data_source_importance_bars,
    create_feature_consistency_plot,
    create_feature_importance_by_grid_level,
)
from entropice.dashboard.utils.loaders import TrainingResult


def _extract_feature_importance_from_results(
    training_results: list[TrainingResult],
) -> pd.DataFrame:
    """Extract feature importance from all training results.

    Args:
        training_results: List of TrainingResult objects

    Returns:
        DataFrame with columns: feature, importance, stddev, model, grid, level, task, target, data_source, grid_level

    """
    fis = []
    for tr in training_results:
        fi = tr.run.feature_importance.reset_index().rename(columns={"index": "feature"})
        fi["model"] = tr.run.model_type
        fi["grid"] = tr.run.dataset.grid
        fi["level"] = tr.run.dataset.level
        fi["task"] = tr.run.task
        fi["target"] = tr.run.target
        fis.append(fi)

    fi = pd.concat(fis, ignore_index=True)
    # Add data source categorization
    fi["data_source"] = fi["feature"].apply(_categorize_feature)
    fi["grid_level"] = fi["grid"] + "_" + fi["level"].astype(str)

    return fi


def _categorize_feature(feature_name: str) -> str:
    """Categorize feature by data source."""
    feature_lower = feature_name.lower()
    if feature_lower.startswith("arcticdem"):
        return "ArcticDEM"
    if feature_lower.startswith("era5"):
        return "ERA5"
    if feature_lower.startswith("embeddings") or feature_lower.startswith("alphaearth"):
        return "Embeddings"
    return "General"


@st.fragment
def render_feature_importance_analysis(training_results: list[TrainingResult]):
    """Render feature importance analysis section.

    Args:
        training_results: List of RandomSearchCV training results

    """
    st.header("🔍 Feature Importance Analysis")

    st.markdown(
        """
        This section analyzes which features are most important across different
        models, grid levels, tasks, and targets.
        """
    )

    # Extract feature importance
    with st.spinner("Extracting feature importance from training results..."):
        fi = _extract_feature_importance_from_results(training_results)

    if fi is None:
        st.warning("No feature importance data available. Model state files may be missing.")
        return

    st.success(f"Extracted feature importance from {len(fi)} feature-model combinations")

    # Filters
    st.subheader("Filters")
    col1, col2, col3 = st.columns(3)

    with col1:
        # Task filter
        available_tasks = ["All", *sorted(fi["task"].unique().tolist())]
        selected_task = st.selectbox("Task", options=available_tasks, index=0, key="fi_task_filter")

    with col2:
        # Target filter
        available_targets = ["All", *sorted(fi["target"].unique().tolist())]
        selected_target = st.selectbox("Target Dataset", options=available_targets, index=0, key="fi_target_filter")

    with col3:
        # Top N features
        top_n_features = st.number_input("Top N Features", min_value=5, max_value=50, value=15, key="top_n_features")

    # Apply filters
    filtered_fi = fi.copy()
    if selected_task != "All":
        filtered_fi = filtered_fi.loc[filtered_fi["task"] == selected_task]
    if selected_target != "All":
        filtered_fi = filtered_fi.loc[filtered_fi["target"] == selected_target]
    if len(filtered_fi) == 0:
        st.warning("No feature importance data available for the selected filters.")
        return

    # Section 1: Top features by grid level
    st.subheader("Top Features by Grid Level")

    try:
        fig = create_feature_importance_by_grid_level(filtered_fi, top_n=top_n_features)
        st.plotly_chart(fig, width="stretch")
    except Exception as e:
        st.error(f"Could not create feature importance by grid level plot: {e}")

    # Show detailed breakdown in expander
    grid_levels = sorted(filtered_fi["grid_level"].unique())

    with st.expander("Show Detailed Breakdown by Grid Level", expanded=False):
        for grid_level in grid_levels:
            grid_data = filtered_fi[filtered_fi["grid_level"] == grid_level]

            # Get top features for this grid level
            top_features_grid = (
                grid_data.groupby("feature")["importance"].mean().reset_index().nlargest(top_n_features, "importance")
            )

            st.markdown(f"**{grid_level.replace('_', '-').title()}**")

            # Create display dataframe with data source
            display_df = top_features_grid.merge(
                grid_data[["feature", "data_source"]].drop_duplicates(), on="feature", how="left"
            )
            display_df.columns = ["Feature", "Mean Importance", "Data Source"]
            display_df = display_df.sort_values("Mean Importance", ascending=False)

            st.dataframe(display_df, width="stretch", hide_index=True)

    # Section 2: Feature importance consistency across models
    st.subheader("Feature Importance Consistency Across Models")

    st.markdown(
        """
        **Coefficient of Variation (CV)**: Lower values indicate more consistent importance across models.
        High CV suggests the feature's importance varies significantly between different models.
        """
    )

    try:
        fig = create_feature_consistency_plot(filtered_fi, top_n=top_n_features)
        st.plotly_chart(fig, width="stretch")
    except Exception as e:
        st.error(f"Could not create feature consistency plot: {e}")

    # Show detailed statistics in expander
    with st.expander("Show Detailed Statistics", expanded=False):
        # Get top features overall
        overall_top_features = (
            filtered_fi.groupby("feature")["importance"]
            .mean()
            .reset_index()
            .nlargest(top_n_features, "importance")["feature"]
            .tolist()
        )

        # Calculate variance in importance across models for each feature
        feature_variance = (
            filtered_fi[filtered_fi["feature"].isin(overall_top_features)]
            .groupby("feature")["importance"]
            .agg(["mean", "std", "min", "max"])
            .reset_index()
        )
        feature_variance["coefficient_of_variation"] = feature_variance["std"] / feature_variance["mean"]
        feature_variance = feature_variance.sort_values("mean", ascending=False)

        # Add data source
        feature_variance = feature_variance.merge(
            filtered_fi[["feature", "data_source"]].drop_duplicates(), on="feature", how="left"
        )

        feature_variance.columns = ["Feature", "Mean", "Std Dev", "Min", "Max", "CV", "Data Source"]

        st.dataframe(
            feature_variance[["Feature", "Data Source", "Mean", "Std Dev", "CV"]],
            width="stretch",
            hide_index=True,
        )

    # Section 3: Feature importance by data source
    st.subheader("Feature Importance by Data Source")

    try:
        fig = create_data_source_importance_bars(filtered_fi)
        st.plotly_chart(fig, width="stretch")
    except Exception as e:
        st.error(f"Could not create data source importance chart: {e}")

    # Show detailed table in expander
    with st.expander("Show Data Source Statistics", expanded=False):
        # Aggregate importance by data source
        source_importance = filtered_fi.groupby("data_source")["importance"].agg(["sum", "mean", "count"]).reset_index()
        source_importance.columns = ["Data Source", "Total Importance", "Mean Importance", "Feature Count"]
        source_importance = source_importance.sort_values("Total Importance", ascending=False)

        st.dataframe(source_importance, width="stretch", hide_index=True)
