"""End-to-end smoke/regression check; no training or uploads."""
import argparse, gc, json, math, platform
from pathlib import Path
import torch
from tiny_jev.inference import TinyJev
from tiny_jev.schema import Question
from tiny_jev.compat import JevAdapter
from tiny_jev import __version__

parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--output", required=True)
args = parser.parse_args()
root = Path(args.model)
assert torch.cuda.is_available(), "Enable a Kaggle GPU for this verification"
model = TinyJev(root, device="cuda:0")
adapter = JevAdapter(model, "tiny-jev-release-verification")
fixtures = json.loads((root / "verification_examples.json").read_text())
maximum = 0.0
for fixture in fixtures:
    questions = {r["id"]: Question.from_record(r) for r in fixture["records"]}
    output = model.decide(fixture["state"], questions)
    specs = {}
    for row in fixture["records"]:
        spec = {"type": row["kind"], "instructions": row["question"]}
        if row["kind"] == "choice":
            spec["criteria"] = dict(zip(row["options"], row["descriptions"]))
        elif row["kind"] == "score":
            spec["criteria"] = row["descriptions"]
        specs[row["id"]] = spec
    adapted = adapter.system_one(state=fixture["state"], questions=specs)
    assert "usage" not in adapted
    for key, answer in adapted["answers"].items():
        assert "confidence" not in answer
        assert answer == {field: output[key][field] for field in answer}
    for r in fixture["records"]:
        a = output[r["id"]]
        p = [1-a["noul"], a["noul"]] if r["kind"] == "noul" else [a["probabilities"][k] for k in r["options"]]
        assert all(math.isfinite(v) and 0 <= v <= 1 for v in p)
        assert abs(sum(p)-1) < 1e-5
        drift = max(abs(x-y) for x,y in zip(p,r["reference_probabilities"]))
        maximum = max(maximum, drift)
        assert drift <= 0.01, f"Unexpected saved-prediction drift: {drift}; investigate before release"
        single = model.decide(fixture["state"], {r["id"]: questions[r["id"]]})[r["id"]]
        q = [1-single["noul"], single["noul"]] if r["kind"] == "noul" else [single["probabilities"][k] for k in r["options"]]
        assert max(abs(x-y) for x,y in zip(p,q)) <= 0.002
    print(json.dumps(output, indent=2))
config = json.loads((root / "config.json").read_text())
report = {"status":"passed", "adapter_status":"passed", "library_version":__version__, "scope":"8 synthetic questions; serving smoke/regression check, not human validation",
          "max_probability_drift":maximum, "allowed_probability_drift":0.01,
          "weights_sha256":config["weights_sha256"], "encoder":config["encoder"],
          "python":platform.python_version(), "torch":str(torch.__version__), "gpu":torch.cuda.get_device_name(0)}
Path(args.output).write_text(json.dumps(report,indent=2))
print("RELEASE VERIFICATION PASSED")