#!/usr/bin/env python3
"""Prepare and evaluate the four registered UNNS Gravity Bridge null branches."""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import numpy as np
import pandas as pd

# Import reusable, preregistered model code without duplicating the estimator.
THIS_DIR = Path(__file__).resolve().parent
if str(THIS_DIR) not in sys.path:
    sys.path.insert(0, str(THIS_DIR))
from run_model_comparison import evaluate_model, feature_sets, load_inputs, load_profiles  # noqa: E402


def read_ladder(path: Path) -> np.ndarray:
    if path.suffix.lower() == ".csv":
        frame = pd.read_csv(path, float_precision="round_trip")
        values = pd.to_numeric(frame.iloc[:, 0], errors="coerce").dropna().to_numpy(dtype=float)
    else:
        values = np.array([float(x) for x in path.read_text(encoding="utf-8").split()], dtype=float)
    values = np.unique(np.sort(values[np.isfinite(values)]))
    if len(values) < 3:
        raise ValueError(f"Invalid ladder {path}")
    return values


def gap_stats(gaps: np.ndarray) -> np.ndarray:
    q1, q3 = np.quantile(gaps, [0.25, 0.75])
    return np.array([gaps.sum(), np.median(gaps), q3 - q1, np.var(gaps)], dtype=float)


def surrogate_gaps(original: np.ndarray, rng: np.random.Generator, attempts: int = 500) -> tuple[np.ndarray, float]:
    """Select a positive random gap vector that approximately preserves four registered summaries."""
    target = gap_stats(original)
    target_scale = np.maximum(np.abs(target), 1e-15)
    best = None
    best_score = float("inf")
    ratio = target[2] / max(target[1], 1e-15)
    base_sigma = float(np.clip(np.log1p(max(ratio, 1e-6)), 0.05, 2.5))
    for _ in range(attempts):
        sigma = float(np.clip(rng.normal(base_sigma, 0.35), 0.03, 3.0))
        gaps = rng.lognormal(mean=0.0, sigma=sigma, size=len(original))
        gaps *= target[0] / gaps.sum()
        score = float(np.mean(((gap_stats(gaps) - target) / target_scale) ** 2))
        if score < best_score:
            best, best_score = gaps, score
    assert best is not None
    return best, best_score


def prepare_n1(root: Path, manifest_path: Path, count: int, seed: int) -> Path:
    manifest = pd.read_csv(manifest_path)
    manifest = manifest[manifest["status"].eq("READY")].copy()
    out_i = root / "ladders/null_inputs/N1_GAP_SURROGATE/struc_i"
    out_p = root / "ladders/null_inputs/N1_GAP_SURROGATE/struc_perc_i"
    out_i.mkdir(parents=True, exist_ok=True)
    out_p.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(seed)
    rows = []
    for m in manifest.itertuples(index=False):
        values = read_ladder(root / str(m.struc_perc_i_input))
        gaps = np.diff(values)
        for s in range(1, count + 1):
            sgaps, score = surrogate_gaps(gaps, rng)
            surrogate = values.min() + np.concatenate([[0.0], np.cumsum(sgaps)])
            sid = f"{m.safe_id}__N1S{s:03d}"
            csv_path = out_i / f"{sid}.csv"
            txt_path = out_p / f"{sid}.txt"
            pd.DataFrame({"log10_g_bar": surrogate}).to_csv(csv_path, index=False, float_format="%.17g")
            txt_path.write_text("\n".join(f"{v:.17g}" for v in surrogate) + "\n", encoding="utf-8")
            rows.append({
                "surrogate_id": sid,
                "base_galaxy_id": m.galaxy_id,
                "surrogate_index": s,
                "fit_score": score,
                "n": len(surrogate),
                "range_relative_error": (np.ptp(surrogate) - np.ptp(values)) / max(np.ptp(values), 1e-15),
                "median_gap_relative_error": (np.median(sgaps) - np.median(gaps)) / max(np.median(gaps), 1e-15),
                "iqr_relative_error": ((np.quantile(sgaps, .75)-np.quantile(sgaps, .25)) - (np.quantile(gaps, .75)-np.quantile(gaps, .25))) / max(np.quantile(gaps, .75)-np.quantile(gaps, .25), 1e-15),
                "variance_relative_error": (np.var(sgaps) - np.var(gaps)) / max(np.var(gaps), 1e-15),
                "struc_i_input": str(csv_path.relative_to(root)),
                "struc_perc_i_input": str(txt_path.relative_to(root)),
            })
    output = root / "ladders/null_inputs/N1_GAP_SURROGATE/n1_surrogate_manifest.csv"
    pd.DataFrame(rows).to_csv(output, index=False)
    return output


def model_rmse(frame: pd.DataFrame, spec: dict, model_name: str = "M_IP") -> float:
    numeric, categorical = feature_sets(spec)[model_name]
    _, _, summary = evaluate_model(
        frame, model_name, numeric, categorical,
        float(spec["estimator"]["alpha"]), int(spec["cross_validation"]["requested_splits"])
    )
    return float(summary["rmse_mean"])


def cyclic_mapping(ids: list[str], shift: int) -> dict[str, str]:
    return {gid: ids[(i + shift) % len(ids)] for i, gid in enumerate(ids)}


def permute_signatures(frame: pd.DataFrame, signature_columns: list[str], rng: np.random.Generator) -> pd.DataFrame:
    galaxy = frame.groupby("galaxy_id", as_index=False).first()
    sort_col = "profile_median_gap" if "profile_median_gap" in galaxy.columns else "galaxy_id"
    ids = galaxy.sort_values(sort_col)["galaxy_id"].astype(str).tolist()
    shift = int(rng.integers(1, len(ids)))
    mapping = cyclic_mapping(ids, shift)
    sig = galaxy.set_index("galaxy_id")[signature_columns]
    out = frame.copy()
    for recipient, donor in mapping.items():
        mask = out["galaxy_id"].eq(recipient)
        out.loc[mask, signature_columns] = sig.loc[donor].to_numpy()
    return out


def permute_target_curves(frame: pd.DataFrame, rng: np.random.Generator) -> pd.DataFrame:
    ids = sorted(frame["galaxy_id"].unique())
    shift = int(rng.integers(1, len(ids)))
    mapping = cyclic_mapping(ids, shift)
    donor_curves = {}
    for gid, g in frame.groupby("galaxy_id"):
        ordered = g.sort_values("r_norm")
        donor_curves[str(gid)] = (ordered["r_norm"].to_numpy(), ordered["target_y"].to_numpy())
    out = frame.copy()
    for recipient, donor in mapping.items():
        mask = out["galaxy_id"].eq(recipient)
        x, y = donor_curves[donor]
        recipient_r = out.loc[mask, "r_norm"].to_numpy()
        out.loc[mask, "target_y"] = np.interp(recipient_r, x, y)
    return out


def evaluate_n1(profiles: pd.DataFrame, null_signatures_path: Path, surrogate_manifest_path: Path, spec: dict) -> list[float]:
    null_sig = pd.read_csv(null_signatures_path)
    sm = pd.read_csv(surrogate_manifest_path)
    null_sig = null_sig.merge(sm[["surrogate_id", "base_galaxy_id", "surrogate_index"]], left_on="galaxy_id", right_on="surrogate_id", how="inner")
    if null_sig.empty:
        raise ValueError("No N1 signatures matched the surrogate manifest")
    chamber_cols = [c for c in null_sig.columns if c.startswith("struc_i_") or c.startswith("perc_")]
    ordinary_cols = ["profile_range", "profile_median_gap", "profile_iqr", "profile_gap_variance"]
    scores = []
    for s, sig_s in null_sig.groupby("surrogate_index"):
        mapped = sig_s.set_index("base_galaxy_id")
        frame = profiles.copy()
        for col in chamber_cols + ordinary_cols:
            if col in mapped.columns:
                frame[col] = frame["galaxy_id"].map(mapped[col])
        scores.append(model_rmse(frame, spec))
    return scores


def summarize_null(null_id: str, scores: list[float], observed: float) -> dict:
    arr = np.asarray(scores, dtype=float)
    if len(arr) == 0:
        return {"null_id": null_id, "status": "NO_RUNS", "runs": 0, "observed_rmse": observed}
    # Lower RMSE is better. p is the fraction of nulls at least as good as observed.
    p = (1 + np.sum(arr <= observed)) / (len(arr) + 1)
    return {
        "null_id": null_id,
        "status": "COMPLETE",
        "runs": len(arr),
        "observed_rmse": observed,
        "null_rmse_mean": float(arr.mean()),
        "null_rmse_sd": float(arr.std(ddof=1)) if len(arr) > 1 else 0.0,
        "null_rmse_min": float(arr.min()),
        "null_rmse_max": float(arr.max()),
        "empirical_p_null_as_good_or_better": float(p),
    }


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[1]
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--mode", choices=["prepare", "evaluate", "all"], default="all")
    p.add_argument("--project-root", type=Path, default=root)
    p.add_argument("--profiles", type=Path, default=root / "data/prepared_profiles/galaxy_profiles.csv")
    p.add_argument("--signatures", type=Path, default=root / "signatures/chamber_signatures.csv")
    p.add_argument("--spec", type=Path, default=root / "config/analysis_specification.json")
    p.add_argument("--null-signatures", type=Path, default=root / "signatures/chamber_signatures_N1.csv")
    p.add_argument("--n1-manifest", type=Path, default=root / "ladders/null_inputs/N1_GAP_SURROGATE/n1_surrogate_manifest.csv")
    p.add_argument("--permutations", type=int, default=None)
    return p.parse_args()


def main() -> int:
    args = parse_args()
    spec = json.loads(args.spec.read_text(encoding="utf-8"))
    seed = int(spec["random_seed"])
    if args.mode in {"prepare", "all"}:
        output = prepare_n1(
            args.project_root,
            args.project_root / "ladders/ladder_manifest.csv",
            int(spec["n1_surrogates_per_galaxy"]),
            seed,
        )
        print(f"Prepared N1 inputs: {output}")
        if args.mode == "prepare":
            return 0

    profiles = load_profiles(args.profiles)
    frame = load_inputs(args.profiles, args.signatures)
    observed = model_rmse(frame, spec)
    n_perm = args.permutations or int(spec["null_permutations"])
    rng = np.random.default_rng(seed)

    results = []
    # N0 is a variability scale, not a permuted predictive score.
    sd_cols = [c for c in frame.columns if c.startswith("struc_i_sd_")]
    n0_values = frame.groupby("galaxy_id")[sd_cols].first().to_numpy(dtype=float).ravel()
    n0_values = n0_values[np.isfinite(n0_values)]
    results.append({
        "null_id": "N0_STRUC_I_REPEAT_VARIABILITY",
        "status": "COMPLETE" if len(n0_values) else "MISSING_REPEAT_DATA",
        "runs": int(len(n0_values)),
        "observed_rmse": observed,
        "null_rmse_mean": None,
        "null_rmse_sd": None,
        "n0_median_feature_sd": float(np.median(n0_values)) if len(n0_values) else None,
        "n0_max_feature_sd": float(np.max(n0_values)) if len(n0_values) else None,
    })

    if args.null_signatures.exists() and args.n1_manifest.exists():
        n1_scores = evaluate_n1(profiles, args.null_signatures, args.n1_manifest, spec)
        results.append(summarize_null("N1_GAP_SURROGATE", n1_scores, observed))
    else:
        results.append({
            "null_id": "N1_GAP_SURROGATE", "status": "PENDING_CHAMBER_RUNS", "runs": 0,
            "observed_rmse": observed,
            "detail": "Run generated N1 ladders through both chambers and build signatures/chamber_signatures_N1.csv"
        })

    sig_cols = [c for c in frame.columns if c.startswith("struc_i_") or c.startswith("perc_")]
    n2_scores = [model_rmse(permute_signatures(frame, sig_cols, rng), spec) for _ in range(n_perm)]
    results.append(summarize_null("N2_CROSS_GALAXY_SIGNATURE", n2_scores, observed))

    n3_scores = [model_rmse(permute_target_curves(frame, rng), spec) for _ in range(n_perm)]
    results.append(summarize_null("N3_GALAXY_TARGET_PERMUTATION", n3_scores, observed))

    out_dir = args.project_root / "outputs/tables"
    out_dir.mkdir(parents=True, exist_ok=True)
    pd.DataFrame(results).to_csv(out_dir / "null_summary.csv", index=False)
    details = {"observed_M_IP_rmse": observed, "N2_scores": n2_scores, "N3_scores": n3_scores}
    (args.project_root / "outputs/reports/null_battery_details.json").write_text(json.dumps(details, indent=2), encoding="utf-8")
    print(pd.DataFrame(results).to_string(index=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
