"""Prove MPS update, checkpoint reload, dynamic ONNX export before full training."""
import json
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import onnxruntime as ort
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
from pii.core import BIO, LABELS, MODEL_REVISION

torch.manual_seed(42)  # Transformers TrainingArguments default seed.
out = Path("artifacts/probe")
out.mkdir(parents=True, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained("artifacts/backbone", local_files_only=True)
model = AutoModelForTokenClassification.from_pretrained("artifacts/backbone", num_labels=len(BIO), id2label=dict(enumerate(BIO)), label2id={v:k for k,v in enumerate(BIO)}, attn_implementation="eager", local_files_only=True).to("mps")
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
report = {"model_revision": MODEL_REVISION, "parameters": sum(p.numel() for p in model.parameters()), "max_length": model.config.max_position_embeddings, "precision":"float32", "device":"mps", "benchmarks":[]}
before = model.classifier.weight.detach().clone()
# Official HF recipe batch=16; measure both short and maximum-length inputs.
for length in (32, model.config.max_position_embeddings):
    x = torch.full((16, length), 100, dtype=torch.long, device="mps")
    x[:,0], x[:,-1] = tokenizer.cls_token_id, tokenizer.sep_token_id
    y = torch.zeros_like(x)
    y[:,0], y[:,-1] = -100, -100
    y[:,1] = 1
    times = []
    for stage in ("warmup", "measured"):
        optimizer.zero_grad(set_to_none=True)
        torch.mps.synchronize()
        start = time.perf_counter()
        loss = model(input_ids=x, attention_mask=torch.ones_like(x), labels=y).loss
        loss.backward()
        finite = all(torch.isfinite(p.grad).all().item() for p in model.parameters() if p.grad is not None)
        optimizer.step()
        torch.mps.synchronize()
        times.append(time.perf_counter()-start)
        assert torch.isfinite(loss).item() and finite
    report["benchmarks"].append({"batch_size":16,"sequence_length":length,"step_seconds":times,"loss":loss.item(),"allocated_bytes":torch.mps.current_allocated_memory(),"driver_bytes":torch.mps.driver_allocated_memory()})
assert not torch.equal(before, model.classifier.weight)
model.eval()
sample = tokenizer("🙂 José Smith: jose@example.com", return_tensors="pt")
sample = {k:v.to("mps") for k,v in sample.items() if k in ("input_ids","attention_mask")}
with torch.no_grad():
    reference = model(**sample).logits.cpu()
start = time.perf_counter()
model.save_pretrained(out / "checkpoint")
torch.save({"optimizer":optimizer.state_dict(),"rng":torch.get_rng_state(),"mps_rng":torch.mps.get_rng_state()}, out / "resume.pt")
report["checkpoint_seconds"] = time.perf_counter()-start
del optimizer
model = model.to("cpu")
reloaded = AutoModelForTokenClassification.from_pretrained(out / "checkpoint", local_files_only=True, attn_implementation="eager").to("mps").eval()
with torch.no_grad():
    actual = reloaded(**sample).logits.cpu()
report["reload_max_abs_difference"] = (actual-reference).abs().max().item()
assert torch.equal(actual,reference)
del reloaded
torch.mps.empty_cache()
tokenizer.save_pretrained(out)
(out / "config.json").write_text(json.dumps({"labels":LABELS,"max_length":model.config.max_position_embeddings,"stride":0,"model_revision":MODEL_REVISION},indent=2))

class Export(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model=model
    def forward(self, input_ids, attention_mask):
        return self.model(input_ids=input_ids,attention_mask=attention_mask).logits

sample = {k:v.cpu() for k,v in sample.items()}
torch.onnx.export(Export(model).eval(), (sample["input_ids"],sample["attention_mask"]), str(out / "model.onnx"), input_names=["input_ids","attention_mask"], output_names=["logits"], dynamic_axes={"input_ids":{0:"batch",1:"sequence"},"attention_mask":{0:"batch",1:"sequence"},"logits":{0:"batch",1:"sequence"}}, opset_version=17, dynamo=False)
session=ort.InferenceSession(str(out / "model.onnx"),providers=["CPUExecutionProvider"])
report["export_checks"]=[]
for text in ("", "Alice Smith", "🙂 José Smith: jose@example.com", "Long ticket. " * 150):
    encoded=tokenizer(text,return_tensors="pt",truncation=True,max_length=model.config.max_position_embeddings)
    encoded={k:v for k,v in encoded.items() if k in ("input_ids","attention_mask")}
    with torch.no_grad():
        logits=model(**encoded).logits.numpy()
    actual=session.run(None,{k:v.numpy() for k,v in encoded.items()})[0]
    report["export_checks"].append({"tokens":encoded["input_ids"].shape[1],"max_abs_logit_difference":float(np.max(np.abs(logits-actual))),"argmax_identical":bool(np.array_equal(logits.argmax(-1),actual.argmax(-1)))})
    assert np.array_equal(logits.argmax(-1),actual.argmax(-1))
(out / "evidence.json").write_text(json.dumps(report,indent=2)+"\n")
print(json.dumps(report,indent=2),flush=True)