# bayesian_noise.py
# ============================================================
# Bayesian noise propagation on RAW X
# Reference location and Δ widths are evaluated at exact clean-fit x50
# Function-only module for notebook use
# ============================================================

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq

from bayesian import load_xy, fit_bayes, P_with, make_trimmed_dataset

# ---------------- Defaults ----------------
XMAX = 5.0
GRID_N = 1000

SIGMA_MULT = 0.129
SIGMA_ADD = 0.144

N_REFIT = 150
N_TTA = 300
SEED = 1234

CI_LEVEL = 0.95
ALPHA = 1.0 - CI_LEVEL
Q_LO, Q_MD, Q_HI = ALPHA / 2.0, 0.5, 1.0 - ALPHA / 2.0

USE_PRIOR_P = True
PRIOR_R = (1.05, 1.05)
TAU = 25.0

plt.rcParams["legend.frameon"] = False
plt.rcParams["axes.titleweight"] = "bold"
plt.rcParams["axes.labelweight"] = "bold"


# ============================================================
# 1) DATA
# ============================================================

def load_original(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
):
    return load_xy(perc=perc, suv_path=suv_path, flags_path=flags_path)


def load_trimmed(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
    value_to_drop=2.48122597,
    tol=1e-3,
):
    X, y = load_xy(perc=perc, suv_path=suv_path, flags_path=flags_path)
    ds = make_trimmed_dataset(X, y, value_to_drop=value_to_drop, tol=tol)
    return ds["X_trim"], ds["y_trim"]


# ============================================================
# 2) FIT + X50
# ============================================================

def fit_once(X, y, rng, use_prior_p=USE_PRIOR_P, prior_r=PRIOR_R, tau=TAU):
    theta, res = fit_bayes(
        X,
        y,
        seed=int(rng.integers(0, 10_000_000)),
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )
    return theta, res


def x_at_p(theta_hat, p_target=0.5, lo=1e-6, hi=6.0, hi_max=100.0):
    """
    Exact root solve for P(AE|x)=p_target.
    """
    f = lambda x: P_with(theta_hat, x) - p_target
    fa, fb = f(lo), f(hi)

    while np.isfinite(fa) and np.isfinite(fb) and fa * fb > 0 and hi < hi_max:
        hi *= 2.0
        fb = f(hi)

    if (not np.isfinite(fa)) or (not np.isfinite(fb)) or fa * fb > 0:
        return np.nan

    try:
        return float(brentq(f, lo, hi))
    except Exception:
        return np.nan


# ============================================================
# 3) NOISE
# ============================================================

def add_noise_mult(X, sigma, rng):
    X = np.asarray(X, float)
    return np.clip(X * np.exp(rng.normal(0, sigma, size=X.shape)), 1e-12, None)


def add_noise_add(X, sigma, rng):
    X = np.asarray(X, float)
    return np.clip(X + rng.normal(0, sigma, size=X.shape), 1e-12, None)


# ============================================================
# 4) BAND HELPERS
# ============================================================

def monotone_nondec(y):
    y = np.asarray(y, float)
    return np.maximum.accumulate(np.clip(y, 0, 1))


def band_quantiles(curves, q_lo=Q_LO, q_md=Q_MD, q_hi=Q_HI):
    C = np.vstack(curves)
    ql = monotone_nondec(np.quantile(C, q_lo, axis=0))
    qm = monotone_nondec(np.quantile(C, q_md, axis=0))
    qh = monotone_nondec(np.quantile(C, q_hi, axis=0))
    return ql, qm, qh


def delta_width_at_x(lo, hi, xc, x0):
    if not np.isfinite(x0):
        return np.nan
    lo_x = float(np.interp(x0, xc, lo))
    hi_x = float(np.interp(x0, xc, hi))
    return hi_x - lo_x


# ============================================================
# 5) CORE ANALYSIS
# ============================================================

def build_bands_for_dataset(
    X,
    y,
    sigma_mult=SIGMA_MULT,
    sigma_add=SIGMA_ADD,
    x_max=XMAX,
    grid_n=GRID_N,
    n_refit=N_REFIT,
    n_tta=N_TTA,
    seed=SEED,
    use_prior_p=USE_PRIOR_P,
    prior_r=PRIOR_R,
    tau=TAU,
):
    rng = np.random.default_rng(seed)
    xc = np.linspace(1e-12, x_max, grid_n)

    theta_clean, res_clean = fit_once(
        X, y, rng,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )
    clean_curve = P_with(theta_clean, xc)
    x50 = x_at_p(theta_clean, p_target=0.5, lo=1e-6, hi=max(6.0, x_max))

    # multiplicative: refit
    curves_refit_m = []
    for _ in range(n_refit):
        Xn = add_noise_mult(X, sigma_mult, rng)
        thetab, resb = fit_once(Xn, y, rng, use_prior_p=use_prior_p, prior_r=prior_r, tau=tau)
        if resb.success and np.isfinite(resb.fun):
            curves_refit_m.append(P_with(thetab, xc))
    if not curves_refit_m:
        raise RuntimeError("No successful multiplicative-refit curves.")
    lo_rm, md_rm, hi_rm = band_quantiles(curves_refit_m)

    # multiplicative: TTA
    curves_tta_m = []
    for _ in range(n_tta):
        xc_n = np.clip(xc * np.exp(rng.normal(0, sigma_mult, size=xc.size)), 1e-12, None)
        curves_tta_m.append(P_with(theta_clean, xc_n))
    lo_tm, md_tm, hi_tm = band_quantiles(curves_tta_m)

    # additive: refit
    curves_refit_a = []
    for _ in range(n_refit):
        Xn = add_noise_add(X, sigma_add, rng)
        thetab, resb = fit_once(Xn, y, rng, use_prior_p=use_prior_p, prior_r=prior_r, tau=tau)
        if resb.success and np.isfinite(resb.fun):
            curves_refit_a.append(P_with(thetab, xc))
    if not curves_refit_a:
        raise RuntimeError("No successful additive-refit curves.")
    lo_ra, md_ra, hi_ra = band_quantiles(curves_refit_a)

    # additive: TTA
    curves_tta_a = []
    for _ in range(n_tta):
        xc_n = np.clip(xc + rng.normal(0, sigma_add, size=xc.size), 1e-12, None)
        curves_tta_a.append(P_with(theta_clean, xc_n))
    lo_ta, md_ta, hi_ta = band_quantiles(curves_tta_a)

    return {
        "xc": xc,
        "theta_clean": theta_clean,
        "fit_success": bool(res_clean.success),
        "clean_curve": clean_curve,
        "x50": x50,
        "mult": {
            "sigma": sigma_mult,
            "refit": (lo_rm, md_rm, hi_rm),
            "tta": (lo_tm, md_tm, hi_tm),
            "D_refit": delta_width_at_x(lo_rm, hi_rm, xc, x50),
            "D_tta": delta_width_at_x(lo_tm, hi_tm, xc, x50),
            "n_refit_success": len(curves_refit_m),
            "n_tta": len(curves_tta_m),
        },
        "add": {
            "sigma": sigma_add,
            "refit": (lo_ra, md_ra, hi_ra),
            "tta": (lo_ta, md_ta, hi_ta),
            "D_refit": delta_width_at_x(lo_ra, hi_ra, xc, x50),
            "D_tta": delta_width_at_x(lo_ta, hi_ta, xc, x50),
            "n_refit_success": len(curves_refit_a),
            "n_tta": len(curves_tta_a),
        },
    }


def run_noise_analysis(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
    value_to_drop=2.48122597,
    tol=1e-3,
    sigma_mult=SIGMA_MULT,
    sigma_add=SIGMA_ADD,
    x_max=XMAX,
    grid_n=GRID_N,
    n_refit=N_REFIT,
    n_tta=N_TTA,
    seed=SEED,
    use_prior_p=USE_PRIOR_P,
    prior_r=PRIOR_R,
    tau=TAU,
):
    X_full, y_full = load_original(
        perc=perc, suv_path=suv_path, flags_path=flags_path
    )
    X_trim, y_trim = load_trimmed(
        perc=perc, suv_path=suv_path, flags_path=flags_path,
        value_to_drop=value_to_drop, tol=tol
    )

    full_res = build_bands_for_dataset(
        X_full, y_full,
        sigma_mult=sigma_mult, sigma_add=sigma_add,
        x_max=x_max, grid_n=grid_n,
        n_refit=n_refit, n_tta=n_tta,
        seed=seed,
        use_prior_p=use_prior_p, prior_r=prior_r, tau=tau,
    )
    trim_res = build_bands_for_dataset(
        X_trim, y_trim,
        sigma_mult=sigma_mult, sigma_add=sigma_add,
        x_max=x_max, grid_n=grid_n,
        n_refit=n_refit, n_tta=n_tta,
        seed=seed + 100,
        use_prior_p=use_prior_p, prior_r=prior_r, tau=tau,
    )

    return {"full": full_res, "trim": trim_res}


# ============================================================
# 6) SUMMARY TABLE
# ============================================================

def make_noise_summary_table(noise_res):
    rows = []
    for dataset_name, block in [("FULL", noise_res["full"]), ("TRIM", noise_res["trim"])]:
        for noise_name, key in [("multiplicative", "mult"), ("additive", "add")]:
            rows.append({
                "Dataset": dataset_name,
                "NoiseType": noise_name,
                "Sigma": float(block[key]["sigma"]),
                "x50": float(block["x50"]),
                "Delta_refit_at_x50": float(block[key]["D_refit"]),
                "Delta_tta_at_x50": float(block[key]["D_tta"]),
                "N_refit_success": int(block[key]["n_refit_success"]),
                "N_tta": int(block[key]["n_tta"]),
            })
    import pandas as pd
    return pd.DataFrame(rows)


# ============================================================
# 7) PLOTTING
# ============================================================

def plot_4panels(
    full_res,
    trim_res,
    save_prefix="noise_cb_raw_4panels_sigma129_0144_x50",
    x_max=XMAX,
    sigma_mult=SIGMA_MULT,
    sigma_add=SIGMA_ADD,
):
    fig, axs = plt.subplots(2, 2, figsize=(12.5, 7.5), dpi=160, sharex=True, sharey=True)

    COL = {"mult": "#1f78b4", "add": "#ff7f00"}
    ALP = {"refit": 0.25, "tta": 0.12}

    def draw_panel(ax, res, mode, letter, dataset_label):
        xc = res["xc"]
        clean = res["clean_curve"]
        x50 = float(res["x50"])

        lo_r, _, hi_r = res[mode]["refit"]
        lo_t, _, hi_t = res[mode]["tta"]
        Dref = res[mode]["D_refit"]
        Dtta = res[mode]["D_tta"]

        ax.plot(xc, clean, color="k", lw=2.6, zorder=3)
        ax.fill_between(xc, lo_r, hi_r, color=COL[mode], alpha=ALP["refit"], zorder=1)
        ax.fill_between(xc, lo_t, hi_t, color=COL[mode], alpha=ALP["tta"], zorder=0)

        # exact x50 reference
        if np.isfinite(x50):
            ax.axvline(x50, color="#666", ls="--", lw=1.2, alpha=0.9, zorder=4)

            # vertical band-width markers at x50
            lo_r_x = float(np.interp(x50, xc, lo_r))
            hi_r_x = float(np.interp(x50, xc, hi_r))
            lo_t_x = float(np.interp(x50, xc, lo_t))
            hi_t_x = float(np.interp(x50, xc, hi_t))

            ax.plot([x50, x50], [lo_r_x, hi_r_x], color=COL[mode], lw=2.2, alpha=0.95, zorder=6)
            ax.scatter([x50, x50], [lo_r_x, hi_r_x], color=COL[mode], s=18, alpha=0.95, zorder=7)

            ax.plot([x50, x50], [lo_t_x, hi_t_x], color=COL[mode], lw=2.2, alpha=0.45, zorder=5)
            ax.scatter([x50, x50], [lo_t_x, hi_t_x], color=COL[mode], s=18, alpha=0.45, zorder=6)

        ax.set_title(letter, fontsize=14, fontweight="bold", pad=8)

        ax.text(
            0.02, 0.96, dataset_label,
            transform=ax.transAxes, ha="left", va="top",
            fontsize=10, color="#111"
        )
        ax.text(
            0.02, 0.88, f"x50={x50:.2f}",
            transform=ax.transAxes, ha="left", va="top",
            fontsize=10, color="#111"
        )
        ax.text(
            0.03, 0.10, f"Δr@x50={Dref:.2f}  Δt@x50={Dtta:.2f}",
            transform=ax.transAxes, ha="left", va="center", fontsize=9,
            bbox=dict(
                facecolor="white",
                edgecolor=COL[mode],
                boxstyle="round,pad=0.25,rounding_size=0.02",
                lw=0.9, alpha=0.95
            )
        )

        ax.set_xlim(0, x_max)
        ax.set_ylim(-0.05, 1.05)
        ax.grid(alpha=0.25)

    draw_panel(axs[0, 0], full_res, "mult", "A", "FULL (F)")
    draw_panel(axs[0, 1], full_res, "add",  "B", "FULL (F)")
    draw_panel(axs[1, 0], trim_res, "mult", "C", "TRIM (T)")
    draw_panel(axs[1, 1], trim_res, "add",  "D", "TRIM (T)")

    axs[1, 0].set_xlabel("X")
    axs[1, 1].set_xlabel("X")
    axs[0, 0].set_ylabel("P(AE | x)")
    axs[1, 0].set_ylabel("P(AE | x)")

    handles = [
        plt.Line2D([0], [0], color="k", lw=2.6, label="clean Bayesian fit"),
        plt.Rectangle((0, 0), 1, 1, facecolor=COL["mult"], alpha=ALP["refit"],
                      label=f"refit band, mult σ={sigma_mult:.3f}"),
        plt.Rectangle((0, 0), 1, 1, facecolor=COL["mult"], alpha=ALP["tta"],
                      label=f"TTA band, mult σ={sigma_mult:.3f}"),
        plt.Rectangle((0, 0), 1, 1, facecolor=COL["add"], alpha=ALP["refit"],
                      label=f"refit band, add σ={sigma_add:.3f}"),
        plt.Rectangle((0, 0), 1, 1, facecolor=COL["add"], alpha=ALP["tta"],
                      label=f"TTA band, add σ={sigma_add:.3f}"),
        plt.Line2D([0], [0], color="#666", lw=1.2, ls="--", label="x50"),
    ]
    fig.legend(handles, [h.get_label() for h in handles],
               loc="lower center", ncol=3, fontsize=10)

    plt.tight_layout(rect=[0, 0.08, 1, 1])

    if save_prefix is not None:
        fig.savefig(f"{save_prefix}.png", dpi=300, bbox_inches="tight")
        fig.savefig(f"{save_prefix}.pdf", bbox_inches="tight")

    return fig, axs