import json, random, sys
import numpy as np
import torch

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
# Union parity: every idea setting is also evaluated for baseline.
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
EPOCHS = 24
NTR, NTE = 400, 200
K, M, DELTA = 8, 24, 2e-3


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    try:
        if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
    except Exception:
        pass


def directions(n, d, rng):
    q = rng.normal(size=(n, d))
    return q / np.maximum(np.linalg.norm(q, axis=1, keepdims=True), 1e-12)


def boundary_value(z):
    # Track solution u=x^3-3xy^2, harmonic in the unit disk.
    return z[..., 0]**3 - 3.0*z[..., 0]*z[..., 1]**2


def wos_paths(x, k, m, seed):
    """Walk on largest interior spheres in the unit disk.

    The source is zero and the terminal value is evaluated after projecting
    near-boundary points to the unit circle. The authoritative supplied
    convention dt=r^2/(2d) is immaterial here because h1=0.
    """
    x = np.asarray(x, dtype=np.float32)
    b, d = x.shape
    rng = np.random.RandomState(int(seed))
    z = np.repeat(x[:, None, :], k, axis=1).copy()
    for _ in range(m):
        norm = np.linalg.norm(z, axis=2)
        r = np.maximum(1.0 - norm, 0.0)
        active = r > DELTA
        u = directions(b*k, d, rng).reshape(b, k, d)
        z = z + r[..., None] * u
        if not np.any(active):
            break
    norm = np.linalg.norm(z, axis=2, keepdims=True)
    z = z / np.maximum(norm, 1e-8)
    return boundary_value(z).astype(np.float32)


def wos_targets(x, k, m, seed):
    return wos_paths(x, k, m, seed).mean(axis=1, keepdims=True)


def math_check():
    # Sphere step identity: E[|z+rU|^2-|z|^2]=r^2.
    rng = np.random.RandomState(2049)
    z = np.array([.2, -.1]); r = .4
    u = directions(200000, 2, rng)
    inc = float(np.mean(np.sum((z+r*u)**2, axis=1)-np.sum(z*z)))
    # The paper's sampling prediction: variance of a K-average is ~1/K.
    x = np.array([[.23, .17]], dtype=np.float32)
    paths = wos_paths(x, 16384, M, 77)[0]
    rows = []
    for k in [1, 2, 4, 8, 16, 32, 64, 128]:
        q = paths[:len(paths)//k*k].reshape(-1, k).mean(axis=1)
        v = float(np.var(q, ddof=1))
        rows.append({"K": k, "variance": v, "K_times_variance": k*v})
    products = np.array([q["K_times_variance"] for q in rows])
    return {"sphere_second_moment_increment": inc, "predicted_increment": r*r,
            "relative_increment_error": abs(inc-r*r)/(r*r),
            "variance_scaling": rows,
            "K_product_cv": float(products.std()/max(products.mean(), 1e-12)),
            "confirmed": bool(abs(inc-r*r)/(r*r) < .02 and products.std()/max(products.mean(), 1e-12) < .12)}


def make_ds(seed, kind):
    d0 = get_dataset("poisson_boundary", seed=seed, n_train=NTR, n_test=NTE)
    ds = {k: d0[k].clone() if torch.is_tensor(d0[k]) else torch.as_tensor(d0[k], dtype=torch.float32)
          for k in ("xtr", "ytr", "xte", "yte")}
    ds.update({"track": "poisson_boundary", "task": "regression", "metric": "mse",
               "input_shape": (2,), "out_dim": 1})
    if kind == "idea":
        ds["ytr"] = torch.as_tensor(wos_targets(d0["xtr"].numpy(), K, M, seed+10000), dtype=torch.float32)
    return ds


def run(kind, lr, seed, return_model=False):
    seed_all(seed)
    ds = make_ds(seed, kind)
    model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"])
    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=float(lr), batch=128,
                                    log=lambda *a, **k: None)
    if net is None:
        raise RuntimeError("bench training failed")
    if return_model:
        return float(metric), net, ds
    return float(metric)


def base_factory(cfg):
    return lambda seed: run("baseline", cfg["lr"], seed)


def idea_factory(cfg):
    return lambda seed: run("idea", cfg["lr"], seed)


def mechanism_signature():
    # Signature uses outputs from trained systems. Prediction: because WOS
    # targets are unbiased harmonic boundary rollouts, larger-K training
    # should not induce a systematic mean shift relative to exact-label MLP.
    rows = []
    for seed in SEEDS:
        bm, bn, bd = run("baseline", 3e-3, seed, True)
        im, inn, idd = run("idea", 3e-3, seed, True)
        with torch.no_grad():
            bdev = next(bn.parameters()).device
            idev = next(inn.parameters()).device
            bp = bn(bd["xte"].to(bdev)).cpu().numpy().reshape(-1)
            ip = inn(idd["xte"].to(idev)).cpu().numpy().reshape(-1)
        delta = float(np.mean(ip-bp))
        rows.append({"seed": seed, "baseline_mse": bm, "idea_mse": im,
                     "trained_prediction_mean_shift_idea_minus_baseline": delta})
    shifts = np.array([r["trained_prediction_mean_shift_idea_minus_baseline"] for r in rows])
    return {"prediction": "WOS stochastic supervision does not create a systematic output mean shift",
            "predicted_shift": 0.0, "observed_shift_mean": float(shifts.mean()),
            "observed_shift_abs_mean": float(np.mean(np.abs(shifts))),
            "trained_systems": rows,
            "confirmed": bool(abs(shifts.mean()) < 0.01)}


def main():
    math_result = math_check()
    base = sweep_baseline(base_factory, GRID, seeds=SWEEP_SEEDS)
    idea_trials = [{"cfg": cfg, "result": evaluate(idea_factory(cfg), SEEDS)} for cfg in GRID]
    best_trial = min(idea_trials, key=lambda q: q["result"]["mean"])
    report = make_report("poisson_boundary", "mlp_tiny", base, best_trial["result"], {
        "idea_config": best_trial["cfg"],
        "idea_sweep": idea_trials,
        "math_check": math_result,
        "mechanism_signature": mechanism_signature(),
        "custom_track": {"name": "poisson_boundary", "file": "bench/custom_tracks/poisson_boundary.py", "domain": "pde"},
        "structural_match": "PDE Dirichlet boundary-value regression on the unit disk; WOS is the intervention."})
    with open("bench_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()