#!/usr/bin/env python3
"""Forge a REAL trained torch surrogate for szl-lambda-gate.
Kernel = ground truth. Surrogate = a tiny torch MLP that predicts the ADVISORY
gate decision `lambda_gate(axes, threshold).passed` — i.e. Λ(axes) >= threshold —
over the canonical 13-axis Yuyay space. The kernel's weighted-geometric-mean Λ
(with non-compensatory zero-routing: any zero/non-finite axis fails the gate) is
the sole labeler; a sample of labels is re-audited by full kernel replay and MUST
agree or the run fails loudly.

Λ IS NOT PROVEN TRUST — it is the ADVISORY weighted geometric mean, uniqueness =
Conjecture 1 (OPEN). The surrogate approximates the gate DECISION, nothing more.
Seeded, receipted, reproducible. Ships .safetensors + config.json."""
import json, os, random, sys, time, hashlib, platform
_here = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if os.path.isdir(os.path.join(_here, "build", "torch-universal")):
    sys.path.insert(0, os.path.join(_here, "build", "torch-universal"))  # in-repo run
else:
    sys.path.insert(0, "/tmp/kernel-probe/szl-lambda-gate/build/torch-universal")  # forge-dev run
import szl_lambda_gate as lg
import numpy as np
import torch
from torch import nn
from safetensors.torch import save_file

SEED = 20260721
random.seed(SEED); np.random.seed(SEED)
torch.manual_seed(SEED)
T0 = time.time()

K = len(lg.YUYAY_AXES)          # 13 canonical axes
THRESHOLD = 0.5
WEIGHTS = lg.yuyay_weights(dtype=torch.float64)   # uniform 1/13, advisory


def kernel_gate(axes_np):
    """Ground truth: lambda_gate(axes).passed for a batch (N,K)."""
    t = torch.tensor(axes_np, dtype=torch.float64)
    res = lg.lambda_gate(t, weights=WEIGHTS, threshold=THRESHOLD)
    return res.passed.numpy().astype(np.int64)


def synth_axes(n):
    """Synthesize axis-score vectors that straddle the gate boundary, including
    non-compensatory zero-route cases (a single zeroed axis must fail)."""
    X = np.random.uniform(0.0, 1.0, size=(n, K)).astype(np.float64)
    # push a chunk toward the boundary region so the label is non-trivial
    hi = np.random.rand(n) < 0.45
    X[hi] = np.random.uniform(0.55, 1.0, size=(hi.sum(), K))
    # inject explicit zero-route rows (one axis exactly 0 -> Λ=0 -> fail)
    zr = np.random.rand(n) < 0.12
    idx = np.where(zr)[0]
    for i in idx:
        X[i, np.random.randint(K)] = 0.0
    return X


class GateMLP(nn.Module):
    def __init__(self, k, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(k, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


# ---- generate (kernel-labeled) ----
N = 40000
Xall = synth_axes(N)
yall = kernel_gate(Xall)

# ground-truth audit: independent fresh kernel replay must agree on a sample
audit_idx = np.array(sorted(random.sample(range(N), 800)))
replay = kernel_gate(Xall[audit_idx])
assert np.array_equal(replay, yall[audit_idx]), "kernel disagrees on audited sample"
audit_checked = int(len(audit_idx))

# split 80/20 deterministic
perm = np.random.permutation(N)
cut = int(N * 0.8)
tr, te = perm[:cut], perm[cut:]
Xtr = torch.tensor(Xall[tr], dtype=torch.float32)
ytr = torch.tensor(yall[tr], dtype=torch.float32)
Xte = torch.tensor(Xall[te], dtype=torch.float32)
yte = torch.tensor(yall[te], dtype=torch.float32)

HIDDEN = 64
model = GateMLP(K, hidden=HIDDEN)
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
lossf = nn.BCEWithLogitsLoss()
EPOCHS = 150
BATCH = 512
model.train()
for ep in range(EPOCHS):
    order = torch.randperm(Xtr.shape[0])
    for b in range(0, Xtr.shape[0], BATCH):
        bi = order[b:b + BATCH]
        opt.zero_grad()
        logits = model(Xtr[bi])
        loss = lossf(logits, ytr[bi])
        loss.backward(); opt.step()

model.eval()
with torch.no_grad():
    pred = (torch.sigmoid(model(Xte)) >= 0.5).long()
    yte_l = yte.long()
    acc = float((pred == yte_l).float().mean())   # fidelity vs kernel
    # per-class recall
    pos = yte_l == 1
    neg = yte_l == 0
    rec_pass = float((pred[pos] == 1).float().mean()) if pos.any() else float("nan")
    rec_fail = float((pred[neg] == 0).float().mean()) if neg.any() else float("nan")

out = os.path.dirname(os.path.abspath(__file__))
state = {k: v.contiguous() for k, v in model.state_dict().items()}
save_file(state, f"{out}/model.safetensors")
model_sha = hashlib.sha256(open(f"{out}/model.safetensors", "rb").read()).hexdigest()

config = {
    "architecture": "GateMLP",
    "task": "advisory-lambda-gate-decision-surrogate",
    "framework": "pytorch",
    "input_dim": K,
    "input_axes": list(lg.YUYAY_AXES),
    "hidden": HIDDEN,
    "layers": ["Linear(13,64)", "ReLU", "Linear(64,64)", "ReLU", "Linear(64,64)", "ReLU", "Linear(64,1)"],
    "output": "logit; sigmoid>=0.5 => predicted gate PASS (Λ>=threshold)",
    "threshold": THRESHOLD,
    "weights": "yuyay uniform 1/13",
    "label_source": "szl_lambda_gate.lambda_gate(axes, weights=yuyay, threshold=0.5).passed",
    "honesty": "predicts the ADVISORY gate DECISION only; Λ is NOT proven trust; uniqueness = Conjecture 1 (open)",
}
with open(f"{out}/config.json", "w") as f:
    json.dump(config, f, indent=2)

n_pos = int(yall.sum()); n_neg = int(N - n_pos)
receipt = {
  "artifact": "SZLHOLDINGS/szl-lambda-gate surrogate v1",
  "role": "advisory Λ gate-decision surrogate (torch MLP) — kernel remains ground truth",
  "generator": {"script": "scripts/forge.py", "seed": SEED, "kernel_version": lg.__version__,
                 "kernel_labelled": True, "kernel_audited_samples": audit_checked,
                 "labeler": "lambda_gate(axes, weights=yuyay_uniform_1/13, threshold=0.5).passed",
                 "axes": list(lg.YUYAY_AXES), "threshold": THRESHOLD},
  "data": {"rows": int(N), "classes": ["GATE_FAIL", "GATE_PASS"],
            "class_counts": {"GATE_FAIL": n_neg, "GATE_PASS": n_pos},
            "split": "80/20 permutation", "features": list(lg.YUYAY_AXES),
            "feature_policy": "13 Yuyay axis scores in [0,1]; includes non-compensatory zero-route rows"},
  "model": {"type": "pytorch GateMLP (3 hidden ReLU layers, 64 units)",
             "params": {"input_dim": K, "hidden": HIDDEN, "epochs": EPOCHS, "batch": BATCH,
                        "lr": 2e-3, "optimizer": "Adam", "loss": "BCEWithLogits", "seed": SEED},
             "file": "model.safetensors", "sha256": model_sha, "config": "config.json"},
  "metrics_MEASURED": {"fidelity_vs_kernel_heldout": round(acc, 4),
                        "test_accuracy": round(acc, 4),
                        "recall_GATE_PASS": round(rec_pass, 4),
                        "recall_GATE_FAIL": round(rec_fail, 4)},
  "environment": {"python": platform.python_version(), "torch": torch.__version__,
                   "numpy": np.__version__, "host": "replit 2-vCPU container",
                   "wall_seconds": round(time.time() - T0, 1)},
  "honesty": "Every number above is MEASURED by this run. Fidelity = agreement%% with the kernel's ADVISORY lambda_gate decision on a held-out split. Λ is the weighted geometric mean, NOT proven trust; uniqueness = Conjecture 1 (open). The surrogate never replaces the kernel gate.",
  "trained_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}
with open(f"{out}/TRAINING_RECEIPT.json", "w") as f:
    json.dump(receipt, f, indent=2)
print(json.dumps(receipt["metrics_MEASURED"], indent=2))
print(f"rows={N} pos={n_pos} neg={n_neg} kernel_audited={audit_checked} wall={receipt['environment']['wall_seconds']}s")