app.py

"""Main Streamlit app for 5G prediction, optimization, and model research workflows."""

import io
from datetime import datetime
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import LabelEncoder

from analysis.correlation_graphs import render_correlation_heatmap
from analysis.explainability import SHAP_AVAILABLE, build_shap_importance_plot, compute_shap_importance
from analysis.feature_selection import run_feature_selection
from analysis.optimization_strategies import generate_optimization_plan
from analysis.research_plots import classification_diagnostics, multiclass_roc_auc, regression_diagnostics
from models.lstm_model import TENSORFLOW_AVAILABLE, train_lstm_regressor
from models.model_registry import evaluate_models, get_model_registry

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
PRETRAINED_DIR = BASE_DIR / "artifacts" / "pretrained"

st.set_page_config(
    page_title="5G ML Optimization System",
    page_icon="📡",
    layout="wide"
)

st.markdown("""
<style>
:root {
    --bg-base: #f2f7ff;
    --surface: #ffffff;
    --ink-strong: #0d1b2a;
    --ink-soft: #375a7f;
    --brand: #0b66c3;
    --ok: #1f8a5b;
    --warn: #ca8a04;
    --bad: #c0392b;
}

.stApp {
    background:
        radial-gradient(circle at 15% 15%, #dbeafe 0%, transparent 38%),
        radial-gradient(circle at 85% 10%, #c7e0ff 0%, transparent 33%),
        linear-gradient(140deg, #f7fbff 0%, #eef6ff 42%, #f8fbff 100%);
}

.big-title {
    font-size: 38px;
    font-weight: 800;
    color: var(--ink-strong);
    letter-spacing: 0.2px;
}

.subtitle {
    font-size: 18px;
    color: var(--ink-soft);
    margin-bottom: 25px;
}

.section-title {
    font-size: 26px;
    font-weight: 700;
    color: #12385f;
    margin-top: 20px;
    margin-bottom: 10px;
}

.recommendation {
    background-color: #0f2747;
    padding: 16px;
    border-left: 6px solid #0b66c3;
    border-radius: 10px;
    color: #e8f1ff;
    margin-bottom: 10px;
}

.kpi-card {
    background-color: var(--surface);
    border: 1px solid #d7e8ff;
    border-radius: 10px;
    padding: 14px;
    box-shadow: 0 8px 26px rgba(16, 57, 104, 0.08);
}

.step-pill {
    display: inline-block;
    background: #dcecff;
    color: #0b3d71;
    border: 1px solid #b9d4f8;
    border-radius: 999px;
    font-size: 12px;
    font-weight: 700;
    letter-spacing: 0.6px;
    text-transform: uppercase;
    padding: 5px 10px;
    margin-bottom: 6px;
}

.status-pill {
    border-radius: 12px;
    color: white;
    padding: 12px 14px;
    margin-bottom: 10px;
    font-weight: 700;
    box-shadow: 0 8px 20px rgba(0, 0, 0, 0.12);
}

.status-ok {
    background: linear-gradient(135deg, #1f8a5b 0%, #1b6f4a 100%);
}

.status-warn {
    background: linear-gradient(135deg, #ca8a04 0%, #a56f03 100%);
}

.status-bad {
    background: linear-gradient(135deg, #c0392b 0%, #9f2d22 100%);
}

.action-card {
    background: #ffffff;
    border: 1px solid #dbe9ff;
    border-left: 6px solid #0b66c3;
    border-radius: 12px;
    padding: 14px;
    margin-bottom: 10px;
    box-shadow: 0 8px 22px rgba(15, 39, 71, 0.07);
}

.action-priority-high {
    border-left-color: #c0392b;
}

.action-priority-medium {
    border-left-color: #ca8a04;
}

.action-priority-low {
    border-left-color: #1f8a5b;
}

.action-priority-info {
    border-left-color: #0b66c3;
}
</style>
""", unsafe_allow_html=True)


@st.cache_resource
def load_pretrained_models():
    return (
        joblib.load(PRETRAINED_DIR / "latency_model.pkl"),
        joblib.load(PRETRAINED_DIR / "throughput_model.pkl"),
        joblib.load(PRETRAINED_DIR / "qos_model.pkl"),
    )


@st.cache_data
def load_default_dataset(path):
    return pd.read_csv(path)


def render_step_label(step_text):
    st.markdown(f'<div class="step-pill">{step_text}</div>', unsafe_allow_html=True)


def apply_plot_theme(fig):
    fig.update_layout(
        template="plotly_white",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(255,255,255,0.92)",
        font={"family": "Trebuchet MS, Verdana, sans-serif", "color": "#163a62"},
        title={"font": {"size": 18, "color": "#13375f"}},
        margin={"l": 20, "r": 20, "t": 50, "b": 20},
    )
    return fig


def render_status_pills(avg_latency, avg_throughput, most_common_qos):
    latency_state = "status-ok"
    latency_text = f"Latency Risk: Healthy ({avg_latency:.1f} ms)"
    if avg_latency >= 70:
        latency_state = "status-bad"
        latency_text = f"Latency Risk: Critical ({avg_latency:.1f} ms)"
    elif avg_latency >= 50:
        latency_state = "status-warn"
        latency_text = f"Latency Risk: Watch ({avg_latency:.1f} ms)"

    throughput_state = "status-ok"
    throughput_text = f"Throughput Health: Strong ({avg_throughput:.1f} Mbps)"
    if avg_throughput < 60:
        throughput_state = "status-bad"
        throughput_text = f"Throughput Health: Weak ({avg_throughput:.1f} Mbps)"
    elif avg_throughput < 100:
        throughput_state = "status-warn"
        throughput_text = f"Throughput Health: Moderate ({avg_throughput:.1f} Mbps)"

    qos_state = "status-ok" if str(most_common_qos).lower() in ["good", "medium"] else "status-bad"
    qos_text = f"QoS Confidence: {most_common_qos}"

    c1, c2, c3 = st.columns(3)
    c1.markdown(f'<div class="status-pill {latency_state}">{latency_text}</div>', unsafe_allow_html=True)
    c2.markdown(f'<div class="status-pill {throughput_state}">{throughput_text}</div>', unsafe_allow_html=True)
    c3.markdown(f'<div class="status-pill {qos_state}">{qos_text}</div>', unsafe_allow_html=True)


def render_strategy_cards(strategy_df, high_only=False):
    display_df = strategy_df.copy()
    if high_only:
        display_df = display_df[display_df["Priority"] == "High"]
        if display_df.empty:
            display_df = strategy_df.copy()

    for _, row in display_df.iterrows():
        priority = str(row["Priority"]).lower()
        priority_class = f"action-priority-{priority}" if priority in ["high", "medium", "low", "info"] else "action-priority-info"
        st.markdown(
            f"""
            <div class="action-card {priority_class}">
                <strong>{row['Priority']} Priority</strong><br/>
                <strong>Issue:</strong> {row['Issue']}<br/>
                <strong>Current:</strong> {row['Current']}<br/>
                <strong>Target:</strong> {row['Target']}<br/>
                <strong>Action:</strong> {row['PreciseAction']}<br/>
                <strong>Estimated Impact:</strong> {row['EstimatedImpact']}
            </div>
            """,
            unsafe_allow_html=True,
        )

st.markdown('<div class="big-title">AI-Driven 5G Network Optimization System</div>', unsafe_allow_html=True)
st.markdown(
    '<div class="subtitle">A Machine Learning prototype for predicting latency, throughput, and QoS classification in 5G networks.</div>',
    unsafe_allow_html=True
)

st.markdown('<div class="section-title">System Purpose</div>', unsafe_allow_html=True)

st.write("""
This application demonstrates how Machine Learning can support smarter 5G network management.

In a real network, conditions change constantly. The number of users, signal quality, packet loss,
mobility, and available bandwidth can all affect the quality of service. Instead of waiting until
the network becomes slow or unstable, this system tries to predict performance in advance.

The prototype takes 5G network parameters as input, predicts latency and throughput, classifies the
network quality, and provides simple optimization recommendations.
""")

predict_tab, ml_lab_tab = st.tabs(["Performance Prediction", "ML Lab: Train, Compare, Select Features"])

st.sidebar.markdown("### Experience Mode")
experience_mode = st.sidebar.radio(
    "Choose your workflow",
    ["Operator Mode", "Research Mode"],
    help="Operator Mode keeps UI concise for field operations. Research Mode enables full diagnostics and explainability.",
)

required_columns = [
    "signal_strength",
    "sinr",
    "connected_users",
    "bandwidth_mhz",
    "packet_loss",
    "jitter",
    "mobility_speed",
]


with predict_tab:
    render_step_label("Step 1 - Input")
    latency_model, throughput_model, qos_model = load_pretrained_models()

    st.sidebar.header("5G Network Parameters")
    input_method = st.sidebar.radio("Choose input method", ["Manual input", "Upload file"], key="predict_input_method")

    uploaded_file = None
    if input_method == "Upload file":
        uploaded_file = st.sidebar.file_uploader("Upload CSV or TXT file", type=["csv", "txt"], key="predict_uploader")
        st.sidebar.info(
            "File must contain these columns: signal_strength, sinr, connected_users, bandwidth_mhz, packet_loss, jitter, mobility_speed"
        )

    if input_method == "Manual input" or uploaded_file is None:
        signal_strength = st.sidebar.slider("Signal Strength (dBm)", -120, -60, -85, key="signal_strength")
        sinr = st.sidebar.slider("SINR", 0, 30, 15, key="sinr")
        connected_users = st.sidebar.slider("Connected Users", 10, 500, 100, key="connected_users")
        bandwidth_mhz = st.sidebar.selectbox("Bandwidth (MHz)", [20, 40, 60, 80, 100], key="bandwidth_mhz")
        packet_loss = st.sidebar.slider("Packet Loss (%)", 0.0, 5.0, 1.0, key="packet_loss")
        jitter = st.sidebar.slider("Jitter (ms)", 1.0, 30.0, 5.0, key="jitter")
        mobility_speed = st.sidebar.slider("Mobility Speed (km/h)", 0.0, 120.0, 20.0, key="mobility_speed")

        input_data = pd.DataFrame({
            "signal_strength": [signal_strength],
            "sinr": [sinr],
            "connected_users": [connected_users],
            "bandwidth_mhz": [bandwidth_mhz],
            "packet_loss": [packet_loss],
            "jitter": [jitter],
            "mobility_speed": [mobility_speed],
        })
    else:
        try:
            input_data = pd.read_csv(uploaded_file)
            missing_columns = [col for col in required_columns if col not in input_data.columns]
            if missing_columns:
                st.error("Uploaded file is missing required columns: " + ", ".join(missing_columns))
                st.stop()
            input_data = input_data[required_columns]
        except Exception:
            st.error("Could not read uploaded file. Please upload a valid CSV/TXT file.")
            st.stop()

    with st.expander("Review Input Summary", expanded=True):
        st.dataframe(input_data, use_container_width=True)

    if st.button("Predict Network Performance", key="predict_button"):
        render_step_label("Step 2 - Prediction")
        latency_predictions = latency_model.predict(input_data)
        throughput_predictions = throughput_model.predict(input_data)
        qos_predictions = qos_model.predict(input_data)

        results = input_data.copy()
        results["predicted_latency_ms"] = latency_predictions
        results["predicted_throughput_mbps"] = throughput_predictions
        results["qos_classification"] = qos_predictions

        st.markdown('<div class="section-title">Prediction Results</div>', unsafe_allow_html=True)

        avg_latency = results["predicted_latency_ms"].mean()
        avg_throughput = results["predicted_throughput_mbps"].mean()
        most_common_qos = results["qos_classification"].mode()[0]

        render_step_label("Step 3 - Status")
        render_status_pills(avg_latency, avg_throughput, most_common_qos)

        col1, col2, col3 = st.columns(3)
        col1.metric("Average Latency", f"{avg_latency:.2f} ms")
        col2.metric("Average Throughput", f"{avg_throughput:.2f} Mbps")
        col3.metric("Most Common QoS", most_common_qos)

        st.dataframe(results, use_container_width=True)

        chart_col1, chart_col2 = st.columns(2)

        with chart_col1:
            latency_fig = go.Figure(go.Indicator(
                mode="gauge+number",
                value=avg_latency,
                title={"text": "Average Latency Level"},
                gauge={
                    "axis": {"range": [0, 120]},
                    "bar": {"color": "#0b66c3"},
                    "steps": [
                        {"range": [0, 40], "color": "#d9f2ff"},
                        {"range": [40, 70], "color": "#fff3cd"},
                        {"range": [70, 120], "color": "#f8d7da"},
                    ],
                },
            ))
            latency_fig = apply_plot_theme(latency_fig)
            st.plotly_chart(latency_fig, use_container_width=True)

        with chart_col2:
            throughput_fig = go.Figure(go.Indicator(
                mode="gauge+number",
                value=avg_throughput,
                title={"text": "Average Throughput Level"},
                gauge={
                    "axis": {"range": [0, 250]},
                    "bar": {"color": "#0b66c3"},
                    "steps": [
                        {"range": [0, 60], "color": "#f8d7da"},
                        {"range": [60, 100], "color": "#fff3cd"},
                        {"range": [100, 250], "color": "#d9f2ff"},
                    ],
                },
            ))
            throughput_fig = apply_plot_theme(throughput_fig)
            st.plotly_chart(throughput_fig, use_container_width=True)

        render_step_label("Step 4 - Action Plan")
        st.markdown('<div class="section-title">Precise Optimization Strategy Plan</div>', unsafe_allow_html=True)
        strategy_df = generate_optimization_plan(input_data, results)
        if experience_mode == "Operator Mode":
            high_only = st.toggle("Show high-priority actions only", value=True, key="strategy_high_only")
            render_strategy_cards(strategy_df, high_only=high_only)
        else:
            render_strategy_cards(strategy_df, high_only=False)
            with st.expander("Strategy table (research view)", expanded=False):
                st.dataframe(strategy_df, use_container_width=True)

        strategy_csv = strategy_df.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="Download Optimization Plan",
            data=strategy_csv,
            file_name="5g_optimization_plan.csv",
            mime="text/csv",
        )

        st.markdown('<div class="section-title">Social and Technical Impact</div>', unsafe_allow_html=True)
        st.write(
            "Reliable 5G networks are important for modern society because they support healthcare systems, "
            "online education, emergency communication, smart cities, industrial automation, autonomous "
            "transportation, and IoT services."
        )

        csv = results.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="Download Prediction Results",
            data=csv,
            file_name="5g_prediction_results.csv",
            mime="text/csv",
        )


with ml_lab_tab:
    render_step_label("Step 1 - Dataset")
    st.markdown('<div class="section-title">Model Training and Feature Engineering Lab</div>', unsafe_allow_html=True)
    st.write(
        "Train and compare multiple models (linear regression, tree models, SVM, KNN, neural networks), "
        "run feature selection, inspect correlations, and optionally benchmark LSTM on regression targets."
    )

    dataset_path = DATA_DIR / "5g_network_data.csv"
    uploaded_dataset = st.file_uploader("Upload a dataset for training (CSV)", type=["csv"], key="training_uploader")

    if uploaded_dataset is not None:
        dataset = pd.read_csv(uploaded_dataset)
    elif dataset_path.exists():
        dataset = load_default_dataset(str(dataset_path))
    else:
        st.warning("No dataset found. Upload a CSV file to continue.")
        st.stop()

    with st.expander("Dataset preview", expanded=True):
        st.dataframe(dataset.head(20), use_container_width=True)

    available_targets = [col for col in dataset.columns if col not in required_columns] or list(dataset.columns)
    default_target = "latency" if "latency" in dataset.columns else available_targets[0]

    col_a, col_b = st.columns(2)
    with col_a:
        target_col = st.selectbox("Target column", options=available_targets, index=available_targets.index(default_target))
    with col_b:
        test_size = st.slider("Test set size", 0.1, 0.4, 0.2, 0.05)

    excluded_cols = {target_col}
    feature_options = [col for col in dataset.columns if col not in excluded_cols]
    default_features = [col for col in required_columns if col in feature_options] or feature_options
    selected_features = st.multiselect("Feature columns", options=feature_options, default=default_features)

    if not selected_features:
        st.warning("Select at least one feature column.")
        st.stop()

    task_type = "classification" if dataset[target_col].dtype == "object" else "regression"
    if dataset[target_col].nunique() <= 8 and dataset[target_col].dtype != "float64":
        task_type = "classification"

    st.info(f"Detected task type: {task_type.title()}")

    render_step_label("Step 2 - Correlation and Feature Design")
    st.markdown("Correlation matrix (numeric columns)")
    corr_fig = render_correlation_heatmap(dataset)
    if corr_fig is not None:
        corr_fig = apply_plot_theme(corr_fig)
        st.plotly_chart(corr_fig, use_container_width=True)
    else:
        st.warning("No numeric columns available to compute a correlation matrix.")

    st.markdown("Feature selection")
    feature_method = st.selectbox("Selection method", ["Mutual Information", "F-Score", "RFE"])
    k_features = st.slider("Number of selected features", 1, len(selected_features), min(4, len(selected_features)))

    st.markdown("Research options")
    default_cv = True if experience_mode == "Research Mode" else False
    run_cv = st.checkbox("Run cross-validation for each model", value=default_cv)
    cv_folds = st.slider("Cross-validation folds", 3, 10, 5) if run_cv else 5
    default_shap = True if experience_mode == "Research Mode" else False
    generate_shap = st.checkbox("Generate SHAP explainability plot (can be slower)", value=default_shap)
    if generate_shap and not SHAP_AVAILABLE:
        st.warning("SHAP package not available. Install dependencies from requirements.txt.")
        generate_shap = False

    use_lstm = False
    lstm_epochs = 25
    if task_type == "regression":
        use_lstm = st.checkbox("Also train an LSTM regressor", value=False)
        if use_lstm:
            if TENSORFLOW_AVAILABLE:
                lstm_epochs = st.slider("LSTM epochs", 10, 100, 25, 5)
            else:
                st.warning("TensorFlow is not installed. Install tensorflow to enable LSTM training.")
                use_lstm = False

    if st.button("Train and Compare Models", key="train_compare"):
        render_step_label("Step 3 - Train and Evaluate")
        train_df = dataset[selected_features + [target_col]].dropna().copy()
        x = train_df[selected_features]
        y = train_df[target_col]

        label_encoder = None
        if task_type == "classification" and y.dtype == "object":
            label_encoder = LabelEncoder()
            y = label_encoder.fit_transform(y)

        stratify_target = y if task_type == "classification" else None
        x_train, x_test, y_train, y_test = train_test_split(
            x,
            y,
            test_size=test_size,
            random_state=42,
            stratify=stratify_target,
        )

        selected_after_fs = run_feature_selection(
            x_train,
            y_train,
            x_columns=selected_features,
            task_type=task_type,
            method=feature_method,
            k_features=k_features,
        )

        x_train = x_train[selected_after_fs]
        x_test = x_test[selected_after_fs]

        st.markdown("Selected features")
        st.write(", ".join(selected_after_fs))

        model_registry = get_model_registry(task_type)
        results_df, best_model_name, best_model = evaluate_models(
            model_registry,
            x_train,
            x_test,
            y_train,
            y_test,
            task_type,
            run_cv=run_cv,
            cv_folds=cv_folds,
        )

        if use_lstm and task_type == "regression" and TENSORFLOW_AVAILABLE:
            lstm_metrics = train_lstm_regressor(x_train, x_test, y_train, y_test, lstm_epochs)
            results_df = pd.concat([
                results_df,
                pd.DataFrame([{
                    "Model": "LSTM",
                    "R2": lstm_metrics["R2"],
                    "MAE": lstm_metrics["MAE"],
                    "RMSE": lstm_metrics["RMSE"],
                }]),
            ], ignore_index=True)
            results_df = results_df.sort_values("R2", ascending=False)

        st.markdown('<div class="section-title">Model Leaderboard</div>', unsafe_allow_html=True)
        st.dataframe(results_df, use_container_width=True)

        leaderboard_csv = results_df.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="Download Leaderboard",
            data=leaderboard_csv,
            file_name=f"leaderboard_{target_col}.csv",
            mime="text/csv",
        )

        metric_col = "R2" if task_type == "regression" else "F1"
        bar_fig = px.bar(
            results_df,
            x="Model",
            y=metric_col,
            color=metric_col,
            title=f"Model Comparison by {metric_col}",
            text_auto=".3f",
        )
        bar_fig = apply_plot_theme(bar_fig)
        st.plotly_chart(bar_fig, use_container_width=True)

        st.markdown('<div class="section-title">Best Classical Model</div>', unsafe_allow_html=True)
        st.write(f"Best model: {best_model_name}")

        y_pred_best = best_model.predict(x_test)
        diagnostics_export_df = None

        estimator = best_model.named_steps.get("model") if isinstance(best_model, Pipeline) else best_model
        if hasattr(estimator, "feature_importances_"):
            importance_df = pd.DataFrame({
                "Feature": selected_after_fs,
                "Importance": estimator.feature_importances_,
            }).sort_values("Importance", ascending=False)
            importance_fig = px.bar(importance_df, x="Feature", y="Importance", title="Feature Importance")
            importance_fig = apply_plot_theme(importance_fig)
            st.plotly_chart(importance_fig, use_container_width=True)
        elif hasattr(estimator, "coef_"):
            coefficients = estimator.coef_[0] if np.ndim(estimator.coef_) > 1 else estimator.coef_
            coef_df = pd.DataFrame({
                "Feature": selected_after_fs,
                "Coefficient": coefficients,
            }).sort_values("Coefficient", key=np.abs, ascending=False)
            coef_fig = px.bar(coef_df, x="Feature", y="Coefficient", title="Model Coefficients")
            coef_fig = apply_plot_theme(coef_fig)
            st.plotly_chart(coef_fig, use_container_width=True)

        show_diagnostics = True if experience_mode == "Research Mode" else st.toggle(
            "Show advanced diagnostics",
            value=False,
            key="operator_show_diagnostics",
        )

        if task_type == "regression" and show_diagnostics:
            st.markdown('<div class="section-title">Regression Diagnostics</div>', unsafe_allow_html=True)
            reg_diag = regression_diagnostics(y_test, y_pred_best)
            diag_col1, diag_col2 = st.columns(2)
            with diag_col1:
                reg_diag["scatter"] = apply_plot_theme(reg_diag["scatter"])
                reg_diag["residual_vs_pred"] = apply_plot_theme(reg_diag["residual_vs_pred"])
                st.plotly_chart(reg_diag["scatter"], use_container_width=True)
                st.plotly_chart(reg_diag["residual_vs_pred"], use_container_width=True)
            with diag_col2:
                reg_diag["residual_hist"] = apply_plot_theme(reg_diag["residual_hist"])
                st.plotly_chart(reg_diag["residual_hist"], use_container_width=True)
                error_quantiles = reg_diag["table"]["Residual"].abs().quantile([0.5, 0.75, 0.9, 0.95]).reset_index()
                error_quantiles.columns = ["Quantile", "AbsoluteResidual"]
                st.dataframe(error_quantiles, use_container_width=True)

            diag_csv = reg_diag["table"].to_csv(index=False).encode("utf-8")
            diagnostics_export_df = reg_diag["table"]
            st.download_button(
                label="Download Regression Diagnostics",
                data=diag_csv,
                file_name=f"diagnostics_{target_col}.csv",
                mime="text/csv",
            )
        elif task_type == "classification" and show_diagnostics:
            st.markdown('<div class="section-title">Classification Diagnostics</div>', unsafe_allow_html=True)
            class_labels = list(label_encoder.transform(label_encoder.classes_)) if label_encoder is not None else sorted(list(np.unique(y_test)))
            class_names = list(label_encoder.classes_) if label_encoder is not None else [str(c) for c in class_labels]
            cls_diag = classification_diagnostics(y_test, y_pred_best, class_labels, display_labels=class_names)
            cls_col1, cls_col2 = st.columns(2)
            with cls_col1:
                cls_diag["confusion"] = apply_plot_theme(cls_diag["confusion"])
                st.plotly_chart(cls_diag["confusion"], use_container_width=True)
            with cls_col2:
                cls_diag["class_accuracy"] = apply_plot_theme(cls_diag["class_accuracy"])
                st.plotly_chart(cls_diag["class_accuracy"], use_container_width=True)

            roc_auc_df = multiclass_roc_auc(best_model, x_test, y_test)
            if roc_auc_df is not None:
                st.markdown("ROC-AUC indicators")
                st.dataframe(roc_auc_df, use_container_width=True)

            cm_csv = cls_diag["confusion_table"].to_csv().encode("utf-8")
            diagnostics_export_df = cls_diag["confusion_table"].copy()
            st.download_button(
                label="Download Confusion Matrix",
                data=cm_csv,
                file_name=f"confusion_matrix_{target_col}.csv",
                mime="text/csv",
            )

        if generate_shap and SHAP_AVAILABLE and show_diagnostics:
            st.markdown('<div class="section-title">Model Explainability (SHAP)</div>', unsafe_allow_html=True)
            shap_importance = compute_shap_importance(best_model, x_test)
            shap_fig = build_shap_importance_plot(shap_importance)
            if shap_fig is not None:
                shap_fig = apply_plot_theme(shap_fig)
                st.plotly_chart(shap_fig, use_container_width=True)
                shap_csv = shap_importance.to_csv(index=False).encode("utf-8")
                st.download_button(
                    label="Download SHAP Importance",
                    data=shap_csv,
                    file_name=f"shap_importance_{target_col}.csv",
                    mime="text/csv",
                )
            else:
                st.warning("SHAP could not generate results for this model. Try a tree-based or linear model.")

        experiment_summary = pd.DataFrame(
            [
                {
                    "Target": target_col,
                    "TaskType": task_type,
                    "FeatureMethod": feature_method,
                    "SelectedFeatureCount": len(selected_after_fs),
                    "RunCrossValidation": run_cv,
                    "CVFolds": cv_folds if run_cv else 0,
                    "BestModel": best_model_name,
                    "TestSize": test_size,
                }
            ]
        )
        with st.expander("Experiment summary", expanded=(experience_mode == "Research Mode")):
            st.dataframe(experiment_summary, use_container_width=True)

        summary_csv = experiment_summary.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="Download Experiment Summary",
            data=summary_csv,
            file_name=f"experiment_summary_{target_col}.csv",
            mime="text/csv",
        )

        if st.button("Save Experiment Package to artifacts/reports", key="save_experiment_package"):
            reports_dir = BASE_DIR / "artifacts" / "reports"
            reports_dir.mkdir(parents=True, exist_ok=True)

            run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
            results_df.to_csv(reports_dir / f"leaderboard_{target_col}_{run_id}.csv", index=False)
            experiment_summary.to_csv(reports_dir / f"summary_{target_col}_{run_id}.csv", index=False)
            pd.DataFrame({"SelectedFeature": selected_after_fs}).to_csv(
                reports_dir / f"selected_features_{target_col}_{run_id}.csv", index=False
            )
            if diagnostics_export_df is not None:
                diagnostics_export_df.to_csv(reports_dir / f"diagnostics_{target_col}_{run_id}.csv", index=False)

            st.success(f"Saved experiment package to {reports_dir}")

        model_bytes = io.BytesIO()
        joblib.dump(best_model, model_bytes)
        st.download_button(
            label="Download Best Classical Model",
            data=model_bytes.getvalue(),
            file_name=f"best_{target_col}_model.pkl",
            mime="application/octet-stream",
        )