#!/usr/bin/env python3
"""Monte Carlo driver for WP28.  One *cell* = (model, p, N, K, DGP options, reps, seed).  Every estimator is computed on the
same simulated sample and the same partition (paired design).  Output: one CSV per cell in results/raw/ with one row per
replication, columns  <estimator>_<j>  for parameter j, plus diagnostics and timings, and a JSON manifest.

Usage
  python3 experiments/run_monte_carlo.py --config experiments/configs/main_exponential.yaml [--cells 0,3,5] [--reps 200]
  python3 experiments/run_monte_carlo.py --list --config ...
"""
from __future__ import annotations
import argparse, hashlib, json, os, pathlib, platform, socket, sys, time
import numpy as np, yaml

ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from pamle.likelihoods import make_model, gather_blocks, STATUS_OK          # noqa: E402
from pamle.algorithms import estimate_all, aggregated_score_hessian, local_split_fits   # noqa: E402
from pamle.partitioning import skewed_sizes                                  # noqa: E402
from pamle import __version__                                                # noqa: E402

SE_ESTIMATORS = ("central", "r1_q0", "r1_q1", "r1_q2", "r2_q0", "r2_q1", "r3_q1", "cs_q1")


def cell_seed(base: int, cell: dict) -> int:
    key = json.dumps({k: cell[k] for k in sorted(cell) if k not in ("reps",)}, sort_keys=True, default=str)
    return (base + int(hashlib.sha256(key.encode()).hexdigest(), 16)) % (2 ** 63)


def expand_cells(cfg: dict):
    cells = []
    for g in cfg["grids"]:
        for N in g["N"]:
            for m in g["m"]:
                if N % m:
                    continue
                for p in g.get("p", [1]):
                    if m < g.get("min_m_over_p", 4) * p:
                        continue
                    c = {"model": g["model"], "p": p, "N": int(N), "K": int(N // m), "m": int(m), "reps": int(g["reps"]),
                         "splits_max": int(g.get("splits_max", 3)), "q_max": int(g.get("q_max", 2)),
                         "dgp": g.get("dgp", {}), "label": g.get("label", g["model"]), "extras": bool(g.get("extras", True))}
                    cells.append(c)
    return cells


def true_theta(cell):
    p = cell["p"]; dgp = cell["dgp"]
    if cell["model"] == "exponential":
        return np.array([float(dgp.get("theta0", 1.0))])
    if "theta0" in dgp:
        return np.asarray(dgp["theta0"], float)
    # default: intercept 0.3, slopes alternating +-0.5 decaying
    th = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(p - 1)])
    if cell["model"] == "poisson":
        th = th * 0.5
    return th


def simulate(model, cell, theta0, rng):
    dgp = cell["dgp"]; N = cell["N"]; p = cell["p"]
    kind = dgp.get("kind", "well_specified")
    if cell["model"] == "exponential":
        x = rng.exponential(1.0 / theta0[0], size=N)
        if kind == "contaminated":
            eps = float(dgp.get("eps", 0.01)); scale = float(dgp.get("scale", 10.0))
            mask = rng.random(N) < eps
            x[mask] = rng.exponential(scale / theta0[0], size=mask.sum())
        return {"x": x}
    cov = None
    if dgp.get("cond", None) is not None and p > 1:      # covariate covariance with prescribed condition number
        q = p - 1; ev = np.geomspace(1.0, 1.0 / float(dgp["cond"]), q)
        cov = np.diag(ev)
    data = model.simulate(rng, N, theta0, cov=cov)
    if kind == "misspecified":                            # omitted quadratic term in the index
        X = data["X"]; eta = X @ theta0 + float(dgp.get("omitted", 0.5)) * (X[:, 1] ** 2 - 1.0)
        if cell["model"] == "logistic":
            data["y"] = (rng.random(N) < 1 / (1 + np.exp(-eta))).astype(float)
        else:
            data["y"] = rng.poisson(np.exp(eta)).astype(float)
    return data


def run_cell(cell, out_dir: pathlib.Path, base_seed: int, reps_override=None, verbose=True):
    model = make_model(cell["model"], cell["p"]); p = model.p
    reps = reps_override or cell["reps"]
    seed = cell_seed(base_seed, cell); rng = np.random.default_rng(seed)
    theta0 = true_theta(cell)
    N, K, m = cell["N"], cell["K"], cell["m"]
    skew = float(cell["dgp"].get("skew", 0.0))
    name = f'{cell["label"]}_p{p}_N{N}_K{K}'
    rows = []; t_start = time.time()
    for b in range(reps):
        data = simulate(model, cell, theta0, rng)
        t0 = time.perf_counter()
        if skew > 1.0:
            # unequal contiguous blocks; estimate_all needs equal shapes -> use K blocks of sizes from skewed_sizes,
            # handled by padding is not possible: we instead run a dedicated weighted/naive comparison (see run_skewed)
            raise NotImplementedError("use run_skewed.py for unequal partitions")
        db = gather_blocks(data, np.arange(N).reshape(K, m))
        out, diag = estimate_all(model, db, splits_max=cell["splits_max"], q_max=cell["q_max"], extras=cell["extras"])
        t_est = time.perf_counter() - t0
        row = {"rep": b, "seed": seed, "N": N, "K": K, "m": m, "p": p, "model": cell["model"], "label": cell["label"],
               "t_estimate_s": t_est, "central_status": diag["central_status"]}
        for s, cnt in diag["local_status_counts"].items():
            row[f"local_fail_s{s}"] = int(sum(cnt[1:]))
        for r in range(1, cell["splits_max"] + 1):
            row[f"valid_r{r}"] = int(diag[f"valid_r{r}"].sum())
            nd = diag[f"newton_r{r}"]
            for q, d in enumerate(nd, start=1):
                row[f"cond_r{r}_q{q}"] = d["cond"]; row[f"step_r{r}_q{q}"] = d["step_norm"]; row[f"grad_r{r}_q{q}"] = d["grad_norm"]
        for k, v in out.items():
            for j in range(p):
                row[f"{k}_{j}"] = float(v[j])
        # standard errors (Fisher at the estimate; sandwich at the centralized estimate)
        for k in SE_ESTIMATORS:
            if k in out and np.all(np.isfinite(out[k])):
                g, H, gk, Hk = aggregated_score_hessian(model, db, out[k])
                try:
                    cov = np.linalg.inv(-H)
                    for j in range(p):
                        row[f"se_{k}_{j}"] = float(np.sqrt(max(cov[j, j], 0) / N))
                except np.linalg.LinAlgError:
                    pass
        th = np.broadcast_to(out["central"], (K, p))
        J = model.outer_scores(th, db).mean(axis=0); g, H, _, _ = aggregated_score_hessian(model, db, out["central"])
        Hi = np.linalg.inv(H); S = Hi @ J @ Hi
        for j in range(p):
            row[f"se_sand_central_{j}"] = float(np.sqrt(max(S[j, j], 0) / N))
        for j in range(p):
            row[f"theta0_{j}"] = float(theta0[j])
        rows.append(row)
        if verbose and (b + 1) % max(1, reps // 10) == 0:
            print(f"[{name}] {b+1}/{reps} reps, {time.time()-t_start:.1f}s", flush=True)
    import pandas as pd
    df = pd.DataFrame(rows)
    out_dir.mkdir(parents=True, exist_ok=True)
    f = out_dir / f"{name}.csv"; df.to_csv(f, index=False)
    sha = hashlib.sha256(f.read_bytes()).hexdigest()
    manifest = {"run_id": name, "cell": cell, "reps": reps, "seed": seed, "pamle_version": __version__,
                "python": platform.python_version(), "numpy": np.__version__, "host": socket.gethostname(),
                "machine": platform.machine(), "wall_seconds": time.time() - t_start, "raw_output_sha256": sha,
                "git_commit": os.popen(f"git -C {ROOT} rev-parse --short HEAD").read().strip(), "status": "completed"}
    (out_dir / f"{name}.manifest.json").write_text(json.dumps(manifest, indent=1))
    return f, manifest


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--cells", default=None, help="comma-separated cell indices (default all)")
    ap.add_argument("--reps", type=int, default=None)
    ap.add_argument("--out", default=str(ROOT / "results" / "raw"))
    ap.add_argument("--list", action="store_true")
    a = ap.parse_args()
    cfg = yaml.safe_load(open(a.config))
    cells = expand_cells(cfg)
    if a.list:
        for i, c in enumerate(cells):
            print(i, c["label"], "p", c["p"], "N", c["N"], "K", c["K"], "m", c["m"], "reps", c["reps"])
        return
    idx = range(len(cells)) if a.cells is None else [int(x) for x in a.cells.split(",")]
    for i in idx:
        f, man = run_cell(cells[i], pathlib.Path(a.out), int(cfg.get("seed_base", 20261028)), a.reps)
        print("wrote", f, f"{man['wall_seconds']:.1f}s")


if __name__ == "__main__":
    main()