"""
OAT (One-At-a-Time) Sensitivity Analysis
=========================================
Analyzes model outputs from OAT parameter sweeps generated by:
  python sensitivity_analysis_param_generation.py --mode oat

For each parameter sweep the following metrics are computed per output:
  - output_range     : max(Y) − min(Y) over the sweep (absolute)
  - pct_change       : output_range / |Y_baseline| × 100
  - direction        : sign of Spearman r between X values and Y values (+1 / −1)
  - signed_si        : direction × pct_change  (Signed Sensitivity Index)

Signed SI is the primary metric used in tornado plots and the heatmap.
A value of +50 means "varying this parameter over its full range increases
the output by 50 % relative to the all-default baseline."

─────────────────────────────────────────────────────────────
HOW TO USE:
  1. Set OAT_PARAMS_FILE  = path to oat_samples.xlsx
  2. Set OAT_OUTPUTS_FILE = path to your OAT model output CSV (row-matched)
  3. Run: python sensitivity_analysis_oat.py
─────────────────────────────────────────────────────────────
"""

import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from scipy import stats
from pathlib import Path

# ─────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────

OAT_PARAMS_FILE  = "oat_runs_manual_bounds.xlsx"
OAT_OUTPUTS_FILE = "../Sex_Trafficking_SDM/OATVariationOutput_manual.csv"

# Outputs to analyze (must exist as columns in your output CSV)
OUTPUTS_OF_INTEREST = [
    "TotalTSS",
    "pctTraffickedVsNontrafficked",
    "Buyers",
    "demand",
    "Traffickers",
    "pctSSCrimRecord",
]

# Human-readable display names for outputs
OUTPUT_LABELS = {
    "TotalTSS":                    "Total Sex Sellers",
    "pctTraffickedVsNontrafficked":"% Trafficked vs Non-Trafficked",
    "Buyers":                      "Buyers",
    "demand":                      "Demand",
    "Traffickers":                 "Traffickers",
    "pctSSCrimRecord":             "% with Criminal Record",
    "TotalSS":                     "Total Sex Sellers (All)",
    "TotalYouthTSS":               "Total Youth Sex Sellers",
    "TotalNonTraffickedSS":        "Non-Trafficked Sex Sellers",
    "EmServicesAdult":             "Adult Em Services",
    "STServicesAdult":             "Adult ST Services",
    "LTServicesAdult":             "Adult LT Services",
    "pctExitedCrimRecord":         "% Exited with Crim Record",
    "cost":                        "Cost",
}

# Max parameters shown in tornado plots and response curve grids; 0 = show all
TOP_N = 15

# Columns of the params file that are not model parameters
NON_PARAM_COLS = ["varied_param", "run_id", "scenario"]

# Directory where figures and CSVs are saved
OUTPUT_DIR = "sensitivity_results_oat"


# ─────────────────────────────────────────────
# DISPLAY HELPERS
# ─────────────────────────────────────────────

def _snake(name: str) -> str:
    """Convert camelCase / PascalCase to snake_case for display."""
    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", name)
    s = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", s)
    return s.lower()


def _out_label(name: str) -> str:
    return OUTPUT_LABELS.get(name, name)


# ─────────────────────────────────────────────
# DATA LOADING
# ─────────────────────────────────────────────

def load_oat_data() -> tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
    """
    Returns (params_df, outputs_df, varied_param_col).

    params_df       : all parameter columns (numeric), index = run_id
    outputs_df      : output columns limited to OUTPUTS_OF_INTEREST that exist
    varied_param_col: string Series identifying which parameter was swept each row
                      ('baseline' for the all-default row)
    """
    params_df  = pd.read_excel(OAT_PARAMS_FILE, sheet_name="oat_params", index_col=0)
    outputs_df = pd.read_csv(OAT_OUTPUTS_FILE, index_col=None)

    print(f"Loaded params : {OAT_PARAMS_FILE}  ({len(params_df)} rows)")
    print(f"Loaded outputs: {OAT_OUTPUTS_FILE}  ({len(outputs_df)} rows)")

    if len(params_df) != len(outputs_df):
        raise ValueError(
            f"Row mismatch: {len(params_df)} param rows vs {len(outputs_df)} output rows.\n"
            "Files must be row-aligned (same run order)."
        )

    if "varied_param" not in params_df.columns:
        raise ValueError(
            "'varied_param' column not found in params file. "
            "Make sure it was generated with --mode oat."
        )

    varied_param = params_df["varied_param"].reset_index(drop=True)
    params_df    = params_df.drop(columns="varied_param").reset_index(drop=True)
    outputs_df   = outputs_df.reset_index(drop=True)

    outputs_present = [c for c in OUTPUTS_OF_INTEREST if c in outputs_df.columns]
    missing = set(OUTPUTS_OF_INTEREST) - set(outputs_present)
    if missing:
        print(f"  WARNING: outputs not found in data: {sorted(missing)}")
    if not outputs_present:
        raise ValueError("No output columns found — check OUTPUTS_OF_INTEREST.")

    return params_df, outputs_df[outputs_present], varied_param


# ─────────────────────────────────────────────
# METRICS
# ─────────────────────────────────────────────

def compute_oat_metrics(
    params_df: pd.DataFrame,
    outputs_df: pd.DataFrame,
    varied_param: pd.Series,
) -> pd.DataFrame:
    """
    For every (parameter, output) pair compute OAT sensitivity metrics.

    Returns a DataFrame indexed by (parameter, output) with columns:
      output_range, pct_change, direction, signed_si,
      y_baseline, param_lo, param_hi, param_default
    """
    # Baseline row
    baseline_mask = varied_param == "baseline"
    if baseline_mask.sum() == 0:
        raise ValueError("No baseline row found (varied_param == 'baseline').")
    baseline_outputs = outputs_df[baseline_mask].iloc[0]
    baseline_params  = params_df[baseline_mask].iloc[0]

    sweep_params = [p for p in varied_param.unique() if p != "baseline"]

    records = []
    for param in sweep_params:
        mask = varied_param == param
        x_vals = params_df.loc[mask, param].values
        sort_idx = np.argsort(x_vals)
        x_sorted = x_vals[sort_idx]

        for out in outputs_df.columns:
            y_vals   = outputs_df.loc[mask, out].values
            y_sorted = y_vals[sort_idx]
            y_base   = baseline_outputs[out]

            out_range = float(y_sorted.max() - y_sorted.min())
            pct_change = (out_range / abs(y_base) * 100) if y_base != 0 else float("nan")

            direction = float(np.sign(stats.spearmanr(x_sorted, y_sorted).statistic))
            signed_si = direction * pct_change

            records.append({
                "parameter":    param,
                "output":       out,
                "output_range": out_range,
                "pct_change":   pct_change,
                "direction":    direction,
                "signed_si":    signed_si,
                "y_baseline":   y_base,
                "param_lo":     float(x_sorted.min()),
                "param_hi":     float(x_sorted.max()),
                "param_default": float(baseline_params.get(param, float("nan"))),
            })

    return pd.DataFrame(records).set_index(["parameter", "output"])


# ─────────────────────────────────────────────
# VISUALIZATION
# ─────────────────────────────────────────────

def tornado_plot_oat(metrics: pd.DataFrame, out_dir: Path):
    """
    One tornado plot per output, parameters ranked by |signed_si|.
    Bar length = signed SI (% change from baseline over full parameter range).
    """
    for out_name in metrics.index.get_level_values("output").unique():
        df = metrics.xs(out_name, level="output")[["signed_si"]].copy()
        df = df.sort_values("signed_si", key=abs, ascending=True)
        if TOP_N > 0:
            df = df.tail(TOP_N)

        fig, ax = plt.subplots(figsize=(9, max(4, len(df) * 0.38)))

        for idx, (param, row) in enumerate(df.iterrows()):
            color = "#2166ac" if row["signed_si"] > 0 else "#d6604d"
            ax.barh(idx, row["signed_si"], color=color, edgecolor="black", linewidth=0.5)

        ax.set_yticks(range(len(df)))
        ax.set_yticklabels([_snake(p) for p in df.index], fontsize=8)
        ax.axvline(0, color="black", linewidth=0.8)
        ax.set_xlabel("Signed Sensitivity Index  (% change from baseline)", fontsize=10)
        ax.set_title(
            f"OAT Sensitivity: {_out_label(out_name)}\n"
            f"Bar = % output change when parameter sweeps full range   "
            f"[blue = positive, red = negative]",
            fontsize=10,
        )
        from matplotlib.patches import Patch
        ax.legend(
            handles=[
                Patch(facecolor="#2166ac", label="Positive effect"),
                Patch(facecolor="#d6604d", label="Negative effect"),
            ],
            loc="lower right", fontsize=8,
        )
        fig.tight_layout()
        path = out_dir / f"oat_tornado_{out_name}.png"
        fig.savefig(path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  Saved: {path.name}")


def heatmap_oat(metrics: pd.DataFrame, out_dir: Path):
    """
    Heatmap of signed SI: rows = parameters, columns = outputs.
    Sorted by strongest effect across any output.
    """
    outputs = metrics.index.get_level_values("output").unique().tolist()
    params  = metrics.index.get_level_values("parameter").unique().tolist()

    matrix = pd.DataFrame(
        {out: metrics.xs(out, level="output")["signed_si"] for out in outputs},
        index=params,
    )
    matrix["_sort"] = matrix.abs().max(axis=1)
    matrix = matrix.sort_values("_sort", ascending=False).drop(columns="_sort")
    if TOP_N > 0:
        matrix = matrix.head(TOP_N)

    vmax = max(0.01, matrix.abs().max().max())
    fig, ax = plt.subplots(figsize=(max(6, len(outputs) * 1.5), max(4, len(matrix) * 0.4)))
    im = ax.imshow(matrix.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
    plt.colorbar(im, ax=ax, label="Signed SI  (% change from baseline)", shrink=0.8)

    ax.set_xticks(range(len(outputs)))
    ax.set_xticklabels([_out_label(o) for o in outputs], rotation=30, ha="right", fontsize=9)
    ax.set_yticks(range(len(matrix)))
    ax.set_yticklabels([_snake(p) for p in matrix.index], fontsize=8)

    for r, param in enumerate(matrix.index):
        for c, out in enumerate(outputs):
            val = matrix.at[param, out]
            txt_color = "white" if abs(val) > vmax * 0.65 else "black"
            ax.text(c, r, f"{val:.0f}", ha="center", va="center", fontsize=7, color=txt_color)

    ax.set_title(
        f"OAT Sensitivity Heatmap — Signed SI (% change from baseline)\n"
        f"Top {min(TOP_N, len(matrix))} parameters by max |SI| across outputs",
        fontsize=10,
    )
    fig.tight_layout()
    path = out_dir / "oat_heatmap.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {path.name}")


def response_curve_plots(
    params_df: pd.DataFrame,
    outputs_df: pd.DataFrame,
    varied_param: pd.Series,
    metrics: pd.DataFrame,
    out_dir: Path,
):
    """
    For each output of interest, plot response curves for the top N most
    influential parameters arranged in a grid. Each subplot shows
    Y vs. X with the baseline Y marked as a horizontal dashed line.
    """
    baseline_mask    = varied_param == "baseline"
    baseline_outputs = outputs_df[baseline_mask].iloc[0]

    for out_name in outputs_df.columns:
        # Rank parameters by |signed_si| for this output
        si_for_out = (
            metrics.xs(out_name, level="output")["signed_si"]
            .abs()
            .sort_values(ascending=False)
        )
        top_params = si_for_out.index.tolist()
        if TOP_N > 0:
            top_params = top_params[:TOP_N]

        n     = len(top_params)
        ncols = min(4, n)
        nrows = (n + ncols - 1) // ncols

        fig = plt.figure(figsize=(5 * ncols, 4 * nrows))
        gs  = gridspec.GridSpec(nrows, ncols, figure=fig, hspace=0.45, wspace=0.35)
        y_base = float(baseline_outputs[out_name])

        for i, param in enumerate(top_params):
            ax   = fig.add_subplot(gs[i // ncols, i % ncols])
            mask = varied_param == param
            x    = params_df.loc[mask, param].values
            y    = outputs_df.loc[mask, out_name].values
            sort_idx = np.argsort(x)

            ax.plot(x[sort_idx], y[sort_idx], color="#2166ac", linewidth=1.6)
            ax.axhline(y_base, color="#d6604d", linewidth=0.9, linestyle="--",
                       label=f"baseline = {y_base:.3g}")
            ax.set_xlabel(_snake(param), fontsize=14)
            ax.set_ylabel(_out_label(out_name), fontsize=14)

            si = metrics.at[(param, out_name), "signed_si"]
            ax.set_title(f"SI = {si:+.1f}%", fontsize=15)
            ax.tick_params(labelsize=13)
            ax.legend(fontsize=13)

        # Hide unused subplots
        for j in range(n, nrows * ncols):
            fig.add_subplot(gs[j // ncols, j % ncols]).set_visible(False)

        fig.suptitle(
            f"OAT Response Curves: {_out_label(out_name)}\n"
            f"Top {n} parameters by |Signed SI|",
            fontsize=17,
        )
        path = out_dir / f"oat_response_{out_name}.png"
        fig.savefig(path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  Saved: {path.name}")


# ─────────────────────────────────────────────
# MAIN
# ─────────────────────────────────────────────

def run_oat_analysis():
    out_dir = Path(OUTPUT_DIR)
    out_dir.mkdir(exist_ok=True)

    print("\n── Loading data ────────────────────────────────────────")
    params_df, outputs_df, varied_param = load_oat_data()
    n_params = varied_param[varied_param != "baseline"].nunique()
    print(f"  Parameters swept : {n_params}")
    print(f"  Outputs found    : {outputs_df.columns.tolist()}")
    print(f"  Total runs       : {len(params_df)}")

    print("\n── Computing OAT metrics ───────────────────────────────")
    metrics = compute_oat_metrics(params_df, outputs_df, varied_param)

    # Save tables
    si_wide = metrics["signed_si"].unstack(level="output")
    si_wide.index = [_snake(p) for p in si_wide.index]
    si_wide.to_csv(out_dir / "oat_signed_si.csv")

    pct_wide = metrics["pct_change"].unstack(level="output")
    pct_wide.index = [_snake(p) for p in pct_wide.index]
    pct_wide.to_csv(out_dir / "oat_pct_change.csv")

    print(f"  Tables saved to '{OUTPUT_DIR}/'")

    # Console summary
    print(f"\n── Top {TOP_N} parameters by |Signed SI| ───────────────────")
    for out_name in outputs_df.columns:
        si = metrics.xs(out_name, level="output")["signed_si"]
        top = si.reindex(si.abs().sort_values(ascending=False).index).head(10)
        print(f"\n  {_out_label(out_name)}:")
        print(f"    {'Parameter':<44} {'Signed SI':>10}  {'% Change':>9}  {'Direction':>10}")
        print(f"    {'─' * 77}")
        for param, si_val in top.items():
            pct = metrics.at[(param, out_name), "pct_change"]
            direction = "positive" if si_val >= 0 else "negative"
            print(f"    {_snake(param):<44} {si_val:>+10.2f}  {pct:>9.2f}%  {direction:>10}")

    print("\n── Generating plots ────────────────────────────────────")
    tornado_plot_oat(metrics, out_dir)
    heatmap_oat(metrics, out_dir)
    response_curve_plots(params_df, outputs_df, varied_param, metrics, out_dir)

    print(f"\nDone. All results saved to '{OUTPUT_DIR}/'")


if __name__ == "__main__":
    run_oat_analysis()
