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

# ----------------------------
# Helpers
# ----------------------------

def safe_read_csv(file):
    df = pd.read_csv(file.name)
    if df.empty:
        raise ValueError("CSV loaded but contains no rows.")
    return df

def infer_defaults(df: pd.DataFrame):
    cols = list(df.columns)

    # turn candidate
    turn_default = "turn" if "turn" in cols else None

    # speaker candidate
    speaker_default = "speaker" if "speaker" in cols else None

    # magnitude candidate
    numeric_cols = list(df.select_dtypes(include=[np.number]).columns)
    mag_default = None
    for c in ["tokens_est", "tokens", "words", "chars", "length"]:
        if c in numeric_cols:
            mag_default = c
            break
    if mag_default is None and numeric_cols:
        mag_default = numeric_cols[0]

    return turn_default, speaker_default, mag_default, cols, numeric_cols

def compute_stability(df: pd.DataFrame, turn_col: str, mag_col: str,
                      rolling_window: int, band_width: float,
                      stability_thresh: float, persistence: int):
    """
    Returns:
      df_out with rolling stats + stable flag
      stable_segments list of (start_turn, end_turn, length)
    """
    d = df.copy()

    if mag_col not in d.columns:
        raise ValueError(f"Selected magnitude column '{mag_col}' not found.")
    if not pd.api.types.is_numeric_dtype(d[mag_col]):
        raise ValueError(f"Selected magnitude column '{mag_col}' is not numeric.")

    # Sort by turn if possible, else keep row order as index
    if turn_col and turn_col in d.columns and pd.api.types.is_numeric_dtype(d[turn_col]):
        d = d.sort_values(turn_col).reset_index(drop=True)
        x = d[turn_col].to_numpy()
    else:
        x = np.arange(len(d))
        turn_col = None  # treat as index

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

    w = int(rolling_window)
    w = max(3, min(w, len(d)))  # enforce sensible range

    s = pd.Series(y)
    roll_mean = s.rolling(w, min_periods=max(3, w//3)).mean().to_numpy()
    roll_std = s.rolling(w, min_periods=max(3, w//3)).std(ddof=0).to_numpy()

    # Avoid division issues
    eps = 1e-9
    z = (y - roll_mean) / (roll_std + eps)

    # Stability criterion:
    # "within stability_thresh sigmas of rolling mean"
    stable = np.abs(z) <= float(stability_thresh)

    # Persistence: stable for N consecutive turns
    p = int(persistence)
    p = max(1, p)

    stable_persist = np.zeros_like(stable, dtype=bool)
    run = 0
    for i, ok in enumerate(stable):
        if ok and not np.isnan(roll_mean[i]) and not np.isnan(roll_std[i]):
            run += 1
        else:
            run = 0
        if run >= p:
            stable_persist[i] = True

    # Bands (for plotting)
    bw = float(band_width)
    upper = roll_mean + bw * roll_std
    lower = roll_mean - bw * roll_std

    d["_x"] = x
    d["_y"] = y
    d["_roll_mean"] = roll_mean
    d["_roll_std"] = roll_std
    d["_band_upper"] = upper
    d["_band_lower"] = lower
    d["_z"] = z
    d["_stable"] = stable_persist

    # Extract stable segments (using stable_persist)
    segments = []
    in_seg = False
    seg_start_idx = None

    for i, ok in enumerate(stable_persist):
        if ok and not in_seg:
            in_seg = True
            seg_start_idx = i
        if in_seg and (not ok or i == len(stable_persist) - 1):
            seg_end_idx = i if ok else i - 1
            in_seg = False

            start_turn = d.loc[seg_start_idx, "_x"]
            end_turn = d.loc[seg_end_idx, "_x"]
            length = seg_end_idx - seg_start_idx + 1
            segments.append((start_turn, end_turn, length))

    return d, segments

def plot_drift_hold(d: pd.DataFrame, title: str, show_points: bool = True):
    fig = plt.figure(figsize=(8, 4.5))
    ax = fig.add_subplot(111)

    ax.plot(d["_x"], d["_roll_mean"], label="Rolling mean")
    ax.plot(d["_x"], d["_band_upper"], label="Band upper")
    ax.plot(d["_x"], d["_band_lower"], label="Band lower")

    if show_points:
        ax.scatter(d["_x"], d["_y"], s=8, alpha=0.6, label="Turns")

    # highlight stable points
    stable_idx = d["_stable"].fillna(False).to_numpy(dtype=bool)
    if stable_idx.any():
        ax.scatter(d.loc[stable_idx, "_x"], d.loc[stable_idx, "_y"], s=14, alpha=0.9, label="Stable (persist)")

    ax.set_title(title)
    ax.set_xlabel("Turn" if "_x" in d.columns else "Index")
    ax.set_ylabel("Magnitude")
    ax.legend()
    fig.tight_layout()
    return fig

# ----------------------------
# Perturbations
# ----------------------------

def temporal_scramble(df: pd.DataFrame, strength: float, seed: int, turn_col: str):
    rng = np.random.default_rng(int(seed))
    d = df.copy()

    if turn_col and turn_col in d.columns and pd.api.types.is_numeric_dtype(d[turn_col]):
        d = d.sort_values(turn_col).reset_index(drop=True)
    else:
        d = d.reset_index(drop=True)

    n = len(d)
    if n < 2 or strength <= 0:
        return d

    window = int(1 + float(strength) * (n - 1))
    window = max(1, min(window, n))

    idx = np.arange(n)
    out = idx.copy()
    for start in range(0, n, window):
        end = min(start + window, n)
        chunk = out[start:end].copy()
        rng.shuffle(chunk)
        out[start:end] = chunk

    return d.iloc[out].reset_index(drop=True)

def metric_noise(df: pd.DataFrame, strength: float, seed: int, col: str):
    rng = np.random.default_rng(int(seed))
    d = df.copy()

    if col not in d.columns:
        raise ValueError(f"Noise column '{col}' not found.")
    if not pd.api.types.is_numeric_dtype(d[col]):
        raise ValueError(f"Noise column '{col}' is not numeric.")

    x = d[col].astype(float).to_numpy()
    if len(x) < 2 or strength <= 0:
        return d

    std = float(np.std(x))
    if std == 0:
        return d

    noise = rng.normal(0, float(strength) * std, size=len(x))
    d[col] = x + noise
    return d

# ----------------------------
# UI callbacks
# ----------------------------

def on_upload(file):
    if file is None:
        return (gr.update(choices=[], value=None),
                gr.update(choices=[], value=None),
                gr.update(choices=[], value=None),
                gr.update(choices=[], value=None),
                gr.update(choices=[], value=None),
                "Upload a CSV to begin.",
                None)

    df = safe_read_csv(file)
    turn_default, speaker_default, mag_default, cols, numeric_cols = infer_defaults(df)

    status = f"Loaded {len(df)} rows, {len(df.columns)} columns."
    preview = df.head(15)

    return (gr.update(choices=cols, value=turn_default),
            gr.update(choices=cols, value=speaker_default),
            gr.update(choices=numeric_cols, value=mag_default),
            gr.update(choices=numeric_cols, value=mag_default),
            gr.update(choices=["None"] + cols, value=turn_default if turn_default else "None"),
            status,
            preview)

def run_drift_hold(file, turn_col, mag_col, rolling_window, band_width, stability_thresh, persistence):
    if file is None:
        return None, "Upload a CSV first.", None

    df = safe_read_csv(file)

    # normalize 'None'
    turn_col = None if (turn_col in [None, "None"] or turn_col not in df.columns) else turn_col

    d, segments = compute_stability(
        df=df,
        turn_col=turn_col,
        mag_col=mag_col,
        rolling_window=int(rolling_window),
        band_width=float(band_width),
        stability_thresh=float(stability_thresh),
        persistence=int(persistence),
    )

    fig = plot_drift_hold(d, title=f"Drift & Hold — {mag_col}")

    if segments:
        seg_lines = "\n".join([f"• Stable segment: {s:.0f} → {e:.0f} (len={L})" for s, e, L in segments[:10]])
        msg = f"Detected {len(segments)} stable segment(s).\n{seg_lines}"
    else:
        msg = "No stable segments detected with current settings."

    return fig, msg, d[["turn"]].head(0) if "turn" not in d.columns else d[["turn", "_y", "_roll_mean", "_roll_std", "_stable"]].head(15)

def run_perturb(file, perturb_type, strength, seed, turn_col_for_scramble, noise_col):
    if file is None:
        return None, "Upload a CSV first.", None

    df = safe_read_csv(file)

    # resolve turn column
    if turn_col_for_scramble in [None, "None"] or turn_col_for_scramble not in df.columns:
        turn_col = None
    else:
        turn_col = turn_col_for_scramble

    if perturb_type == "Temporal scramble":
        df2 = temporal_scramble(df, strength=float(strength), seed=int(seed), turn_col=turn_col)
        msg = f"Applied temporal scramble (strength={strength})."
    else:
        df2 = metric_noise(df, strength=float(strength), seed=int(seed), col=noise_col)
        msg = f"Applied metric noise to '{noise_col}' (strength={strength})."

    # quick default plot column
    _, _, mag_default, _, numeric_cols = infer_defaults(df2)
    if not mag_default:
        return None, "No numeric columns available to plot.", df2.head(15)

    # x-axis
    if turn_col and pd.api.types.is_numeric_dtype(df2[turn_col]):
        x = df2[turn_col].to_numpy()
    else:
        x = np.arange(len(df2))

    fig = plt.figure(figsize=(8, 4.5))
    ax = fig.add_subplot(111)
    ax.plot(x, df2[mag_default].astype(float).to_numpy())
    ax.set_title(f"Perturbed — {mag_default}")
    ax.set_xlabel(turn_col if turn_col else "Index")
    ax.set_ylabel(mag_default)
    fig.tight_layout()

    return fig, msg, df2.head(15)

# ----------------------------
# App
# ----------------------------

with gr.Blocks(title="Threadscope: Drift & Hold") as demo:
    gr.Markdown(
        "Threadscope: Drift & Hold — Bring Your Own Thread\n\n"
        "Upload a CSV to visualize long-form interaction dynamics. "
        "Processed in-session only (no storage)."
    )

    with gr.Row():
        file = gr.File(label="Upload CSV", file_types=[".csv"])

    status = gr.Textbox(label="Status", interactive=False)

    with gr.Row():
        turn_col = gr.Dropdown(label="Turn column (optional)", choices=[], value=None)
        speaker_col = gr.Dropdown(label="Speaker column (optional)", choices=[], value=None)
        mag_col = gr.Dropdown(label="Magnitude column (numeric)", choices=[], value=None)

    with gr.Row():
        preview = gr.Dataframe(label="Preview (first 15 rows)", interactive=False, wrap=True)

    file.change(
        fn=on_upload,
        inputs=[file],
        outputs=[turn_col, speaker_col, mag_col,  # main mapping
                 # reuse: noise column + scramble turn dropdown will be set by same list later
                 # We'll set them in the UI below using same choices/values:
                 # placeholders:
                 ],
    )

    # Workaround: Gradio requires explicit outputs; we'll update extra dropdowns via a second handler
    noise_col = gr.Dropdown(label="Noise column (numeric)", choices=[], value=None)
    turn_col_for_scramble = gr.Dropdown(label="Turn column for scramble (optional)", choices=[], value="None")

    # Update all dropdowns + status + preview on upload
    def on_upload_all(file):
        tc, sc, mc, nc, tcs, st, pv = on_upload(file)
        return tc, sc, mc, nc, tcs, st, pv

    file.change(
        fn=on_upload_all,
        inputs=[file],
        outputs=[turn_col, speaker_col, mag_col, noise_col, turn_col_for_scramble, status, preview],
    )

    with gr.Tabs():
        with gr.Tab("Drift & Hold"):
            with gr.Row():
                rolling_window = gr.Slider(3, 200, value=25, step=1, label="Rolling window (turns)")
                band_width = gr.Slider(0.5, 4.0, value=2.0, step=0.1, label="Band width (σ multiplier)")
            with gr.Row():
                stability_thresh = gr.Slider(0.5, 4.0, value=1.0, step=0.1, label="Stability threshold (|z| ≤ σ)")
                persistence = gr.Slider(1, 100, value=10, step=1, label="Required persistence (turns)")

            run_btn = gr.Button("Run Drift & Hold")
            out_plot = gr.Plot(label="Drift & Hold plot")
            out_msg = gr.Textbox(label="Notes", interactive=False)
            out_table = gr.Dataframe(label="Computed preview", interactive=False, wrap=True)

            run_btn.click(
                fn=run_drift_hold,
                inputs=[file, turn_col, mag_col, rolling_window, band_width, stability_thresh, persistence],
                outputs=[out_plot, out_msg, out_table],
            )

        with gr.Tab("Perturbations"):
            perturb_type = gr.Dropdown(
                label="Perturbation type",
                choices=["Temporal scramble", "Metric noise injection"],
                value="Temporal scramble",
            )
            with gr.Row():
                strength = gr.Slider(0, 1, value=0.35, step=0.01, label="Strength")
                seed = gr.Number(value=7, precision=0, label="Seed")

            run_perturb_btn = gr.Button("Apply perturbation")
            pert_plot = gr.Plot(label="Perturbed plot")
            pert_msg = gr.Textbox(label="Notes", interactive=False)
            pert_preview = gr.Dataframe(label="Perturbed preview (first 15 rows)", interactive=False, wrap=True)

            run_perturb_btn.click(
                fn=run_perturb,
                inputs=[file, perturb_type, strength, seed, turn_col_for_scramble, noise_col],
                outputs=[pert_plot, pert_msg, pert_preview],
            )

demo.launch()