"""Regression Analysis Section."""

import streamlit as st

from entropice.dashboard.plots.regression import plot_regression_scatter, plot_residuals
from entropice.dashboard.utils.loaders import TrainingResult
from entropice.ml.dataset import DatasetEnsemble


def render_regression_analysis(selected_result: TrainingResult):
    """Render regression analysis with true vs predicted scatter plots.

    Args:
        selected_result: The selected TrainingResult object.

    """
    st.header("📊 Regression Analysis")

    # Check if this is a regression task
    if selected_result.run.task in ["binary", "count_regimes", "density_regimes"]:
        st.info("📈 Regression analysis is only available for regression tasks (count, density).")
        return

    # Load predictions
    predictions_df = selected_result.load_predictions()
    if predictions_df is None:
        st.warning("⚠️ No prediction data found for this training result.")
        return

    # Create DatasetEnsemble from settings
    with st.spinner("Loading training data to get true values..."):
        ensemble = DatasetEnsemble(
            grid=selected_result.run.dataset.grid,
            level=selected_result.run.dataset.level,
            members=selected_result.run.dataset.members,
            temporal_mode=selected_result.run.dataset.temporal_mode,
            dimension_filters=selected_result.run.dataset.dimension_filters,
            variable_filters=selected_result.run.dataset.variable_filters,
        )

        # Create training set to get true values
        training_set = ensemble.create_training_set(
            task=selected_result.run.task,
            target=selected_result.run.target,
            device="cpu",
            cache_mode="read",
        )

    # Get split information
    split_series = training_set.split

    # Merge predictions with true values and split info
    # predictions_df should have 'cell_id' and 'predicted' columns
    # training_set.targets has 'y' (true values) with cell_id as index
    true_values = training_set.targets[["y"]].reset_index()

    # Merge on cell_id
    merged = predictions_df.merge(true_values, on="cell_id", how="inner")
    merged["split"] = split_series.reindex(merged["cell_id"]).values

    # Get train, test, and complete data
    train_data = merged[merged["split"] == "train"]
    test_data = merged[merged["split"] == "test"]

    if len(train_data) == 0 or len(test_data) == 0:
        st.error("❌ Could not properly split data into train and test sets.")
        return

    # Display scatter plots
    st.subheader("True vs Predicted Values")
    st.caption("Scatter plots showing the relationship between true and predicted values")

    cols = st.columns(3)

    with cols[0]:
        st.markdown("#### Test Set")
        st.caption("Held-out test set")
        fig_test = plot_regression_scatter(
            test_data["y"],
            test_data["predicted"],
            title="Test Set",
        )
        st.plotly_chart(fig_test, use_container_width=True)

    with cols[1]:
        st.markdown("#### Training Set")
        st.caption("Training set")
        fig_train = plot_regression_scatter(
            train_data["y"],
            train_data["predicted"],
            title="Training Set",
        )
        st.plotly_chart(fig_train, use_container_width=True)

    with cols[2]:
        st.markdown("#### Combplete")
        st.caption("Train + Test sets")
        fig_complete = plot_regression_scatter(
            merged["y"],
            merged["predicted"],
            title="Complete",
        )
        st.plotly_chart(fig_complete, use_container_width=True)

    # Display residual plots
    st.subheader("Residual Analysis")
    st.caption("Residual plots to assess model fit and identify patterns in errors")

    cols = st.columns(3)

    with cols[0]:
        fig_test_res = plot_residuals(test_data["y"], test_data["predicted"], title="Test Set Residuals")
        st.plotly_chart(fig_test_res, use_container_width=True)

    with cols[1]:
        fig_train_res = plot_residuals(train_data["y"], train_data["predicted"], title="Training Set Residuals")
        st.plotly_chart(fig_train_res, use_container_width=True)

    with cols[2]:
        fig_complete_res = plot_residuals(merged["y"], merged["predicted"], title="Complete Residuals")
        st.plotly_chart(fig_complete_res, use_container_width=True)
