import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# --- QUICK VIEW CALLBACK (new, lightweight) ---
def run_quick_view(file, turn_col, mag_col, rolling_window):
    if file is None:
        return None, "Upload a CSV first."

    df = pd.read_csv(file.name)
    if df.empty:
        return None, "CSV is empty."

    # Resolve turn axis
    use_turn = (turn_col not in [None, "None"] and turn_col in df.columns and pd.api.types.is_numeric_dtype(df[turn_col]))
    if use_turn:
        d = df.sort_values(turn_col).reset_index(drop=True)
        x = d[turn_col].to_numpy()
    else:
        d = df.reset_index(drop=True)
        x = np.arange(len(d)) + 1  # friendlier 1-based index

    # Validate magnitude
    if mag_col not in d.columns or not pd.api.types.is_numeric_dtype(d[mag_col]):
        return None, f"'{mag_col}' is not a numeric magnitude column."

    y = d[mag_col].astype(float).to_numpy()

    w = int(rolling_window)
    w = max(3, min(w, len(d)))
    roll = pd.Series(y).rolling(w, min_periods=max(3, w // 3)).mean().to_numpy()

    fig = plt.figure(figsize=(8, 4.5))
    ax = fig.add_subplot(111)
    ax.scatter(x, y, s=8, alpha=0.6, label="Turns")
    ax.plot(x, roll, label=f"Rolling mean (w={w})")
    ax.set_title(f"Quick View — {mag_col}")
    ax.set_xlabel(turn_col if use_turn else "Turn (row order)")
    ax.set_ylabel("Magnitude")
    ax.legend()
    fig.tight_layout()

    msg = "Quick View: response magnitude over time. Use Advanced tabs for stability bands + perturbations."
    return fig, msg

# --- UPLOAD CALLBACK ---
def on_upload_all(file):
    if file is None:
        # Ensure 'None' is always a possible choice for value=None updates
        return (
            gr.Dropdown.update(choices=[None], value=None),
            gr.Dropdown.update(choices=[None], value=None),
            gr.Dropdown.update(choices=[None], value=None),
            gr.Dropdown.update(choices=[None], value=None),
            gr.Dropdown.update(choices=["None"], value="None"),
            "Please upload a CSV file.",
            pd.DataFrame(),
        )

    try:
        df = pd.read_csv(file.name)
        if df.empty:
            return (
                gr.Dropdown.update(choices=[None], value=None),
                gr.Dropdown.update(choices=[None], value=None),
                gr.Dropdown.update(choices=[None], value=None),
                gr.Dropdown.update(choices=[None], value=None),
                gr.Dropdown.update(choices=["None"], value="None"),
                "CSV is empty.",
                pd.DataFrame(),
            )

        columns = list(df.columns)
        numeric_cols = [col for col in columns if pd.api.types.is_numeric_dtype(df[col])]

        # Ensure 'None' is a selectable option for optional columns where value is None
        dropdown_choices_for_none_value = [None] + columns
        dropdown_choices_for_numeric_none_value = [None] + numeric_cols
        dropdown_choices_for_string_none_value = ["None"] + columns

        # Set default mag_col to the first numeric column if available
        default_mag_col = numeric_cols[0] if numeric_cols else None

        return (
            gr.Dropdown.update(choices=dropdown_choices_for_none_value, value=None),
            gr.Dropdown.update(choices=dropdown_choices_for_none_value, value=None),
            gr.Dropdown.update(choices=dropdown_choices_for_numeric_none_value, value=default_mag_col),
            gr.Dropdown.update(choices=dropdown_choices_for_numeric_none_value, value=None),
            gr.Dropdown.update(choices=dropdown_choices_for_string_none_value, value="None"),
            f