"""CPU dynamic INT8 experiment on the reference model only."""
from pathlib import Path
import copy
import json
import sys
import time
import warnings
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from safetensors.torch import load_file
from nexora.model import NexoraLM, ModelConfig
from nexora.training import batch
from nexora.evaluation import percentiles


def main():
    torch.set_num_threads(4)
    cfg = ModelConfig(**json.loads(Path("artifacts/tiny/config.json").read_text()))
    model = NexoraLM(cfg).eval()
    model.load_state_dict(load_file("artifacts/tiny/model.safetensors"))
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", DeprecationWarning)
        quant = torch.ao.quantization.quantize_dynamic(copy.deepcopy(model), {torch.nn.Linear}, dtype=torch.qint8)
    ids = torch.from_numpy(np.load("artifacts/data/validation.npy").astype(np.int64))
    x, y = batch(ids, 4, 128, torch.Generator().manual_seed(9), "cpu")
    results = {}
    logits = {}
    for name, m in [("fp32", model), ("dynamic_int8_linear_only", quant)]:
        times = []
        with torch.no_grad():
            m(x, y)
            for _ in range(20):
                start = time.perf_counter()
                output, loss = m(x, y)
                times.append(time.perf_counter()-start)
        logits[name] = output
        results[name] = {"validation_loss": loss.item(), "batch_latency_seconds": percentiles(times), "samples": len(times)}
    results["max_logit_difference"] = (logits["fp32"]-logits["dynamic_int8_linear_only"]).abs().max().item()
    results["argmax_agreement"] = (logits["fp32"].argmax(-1) == logits["dynamic_int8_linear_only"].argmax(-1)).float().mean().item()
    results["limitations"] = "Tiny model, one validation batch, CPU linear-layer INT8 only. Not INT8 embeddings, BF16/FP8/INT4 deployment, coding/reasoning/context/tool quality validation. Do not infer production speedup."
    Path("reports/quantization.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results))


if __name__ == "__main__":
    main()