#!/usr/bin/env python3
"""Offline GPU training/export only. Deployed inference and bridges are Go."""

import argparse
import copy
import hashlib
import json
import math
import random
import resource
import struct
import sys
import time
from pathlib import Path

import numpy as np
import torch
from torch import nn

LABELS = ["add", "subtract", "multiply", "less_than", "less_equal", "equal", "and", "or"]
FEATURES, HIDDEN, LIMIT = 256, 48, 512
FROZEN_DATASET_SHA = "af0a637320a8256cd6ebcdb3486a6885f2766441c90f0cb154c11599f3c8ca3d"


def require(condition, message):
    if not condition:
        raise ValueError(message)


def unique_pairs(pairs):
    result = {}
    for key, value in pairs:
        require(key not in result, "duplicate JSON key")
        result[key] = value
    return result


def validate_frozen_dataset(path):
    root = Path(__file__).resolve().parents[1]
    raw = path.read_bytes()
    manifest_raw = path.with_name("manifest.json").read_bytes()
    manifest = json.loads(manifest_raw, object_pairs_hook=unique_pairs)
    require(sha(raw) == FROZEN_DATASET_SHA == manifest.get("dataset_sha256"), "dataset differs from frozen pilot")
    require(manifest.get("model_contract_sha256") == sha((root / "model-contract.json").read_bytes()), "model contract pin mismatch")
    require(manifest.get("generator_source_sha256") == sha((root / "cmd/dataset/main.go").read_bytes()), "generator source pin mismatch")
    lines = raw.decode("utf-8").splitlines(keepends=True)
    rows = [json.loads(line, object_pairs_hook=unique_pairs) for line in lines]
    fields = {"id", "template_id", "configuration_id", "language", "split", "text", "label"}
    require(all(isinstance(r, dict) and set(r) == fields for r in rows), "dataset row shape mismatch")
    require(all(isinstance(r["text"], str) and 1 <= len(r["text"].encode("utf-8")) <= LIMIT for r in rows), "input byte limit")
    require(all(r["label"] in LABELS and r["split"] in ("train", "calibration", "test") and r["language"] in ("en", "ko") for r in rows), "dataset enum mismatch")
    require(len({r["id"] for r in rows}) == len(rows) == 2048, "duplicate IDs or wrong row count")
    require(len({r["text"] for r in rows}) == len(rows), "duplicate texts")
    require(manifest.get("rows") == {"total": 2048, "train": 1536, "calibration": 256, "test": 256}, "manifest count mismatch")
    split_rows = {s: [r for r in rows if r["split"] == s] for s in ("train", "calibration", "test")}
    for split, expected in (("train", 1536), ("calibration", 256), ("test", 256)):
        require(len(split_rows[split]) == expected, "split count mismatch")
        split_raw = "".join(line for line, row in zip(lines, rows) if row["split"] == split).encode("utf-8")
        require(sha(split_raw) == manifest.get("split_sha256", {}).get(split), "split hash mismatch")
    for a, b in (("train", "calibration"), ("train", "test"), ("calibration", "test")):
        require(not {r["template_id"] for r in split_rows[a]} & {r["template_id"] for r in split_rows[b]}, "template split leakage")
    return raw, rows, {"dataset_manifest_sha256": sha(manifest_raw), "model_contract_sha256": manifest["model_contract_sha256"],
                       "generator_source_sha256": manifest["generator_source_sha256"]}


def sha(raw):
    return hashlib.sha256(raw).hexdigest()


def save_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True, allow_nan=False) + "\n")


def features(text):
    encoded = text.encode("utf-8")
    if not 1 <= len(encoded) <= LIMIT:
        raise ValueError("input must contain 1..512 UTF-8 bytes; inputs are not truncated")
    raw = bytes(c + 32 if 65 <= c <= 90 else c for c in encoded)
    counts = np.zeros(FEATURES, dtype=np.float32)
    for n in (2, 3):
        for start in range(max(0, len(raw) - n + 1)):
            value = 2166136261
            for c in raw[start:start + n]:
                value = ((value ^ c) * 16777619) & 0xffffffff
            counts[value % FEATURES] += 1
    norm = math.sqrt(sum(float(c) ** 2 for c in counts))
    if norm:
        counts *= np.float32(1 / norm)
    return counts


def ternary(weight):
    scale = weight.detach().abs().mean().clamp_min(1e-8)
    codes = torch.round(weight / scale).clamp(-1, 1)
    quantized = codes * scale
    return weight + (quantized - weight).detach()


class Model(nn.Module):
    def __init__(self, qat=False):
        super().__init__()
        self.first = nn.Linear(FEATURES, HIDDEN)
        self.last = nn.Linear(HIDDEN, len(LABELS))
        self.qat = qat

    def forward(self, x):
        w1 = ternary(self.first.weight) if self.qat else self.first.weight
        w2 = ternary(self.last.weight) if self.qat else self.last.weight
        x = torch.relu(nn.functional.linear(x, w1, self.first.bias))
        return nn.functional.linear(x, w2, self.last.bias)


def probabilities(logits, temperature):
    z = logits.astype(np.float64) / temperature
    z -= z.max(axis=1, keepdims=True)
    p = np.exp(z)
    return p / p.sum(axis=1, keepdims=True)


def scores(logits, labels, temperature):
    p = probabilities(logits, temperature)
    choices = p.argmax(axis=1)
    correct = choices == labels
    onehot = np.eye(len(LABELS))[labels]
    confidence = p.max(axis=1)
    ece, bins = 0.0, []
    for i in range(10):
        mask = (confidence >= i / 10) & (confidence < (i + 1) / 10 if i < 9 else confidence <= 1)
        n = int(mask.sum())
        accuracy = float(correct[mask].mean()) if n else None
        mean = float(confidence[mask].mean()) if n else None
        if n:
            ece += n / len(labels) * abs(accuracy - mean)
        bins.append({"lower": i / 10, "upper": (i + 1) / 10, "count": n,
                     "accuracy": accuracy, "mean_confidence": mean})
    return {"planned": len(labels), "observed": len(labels), "correct": int(correct.sum()),
            "accuracy": float(correct.mean()), "nll": float(-np.log(p[np.arange(len(labels)), labels].clip(1e-12)).mean()),
            "brier_multiclass_sum": float(((p - onehot) ** 2).sum(axis=1).mean()),
            "ece_equal_width_10": ece, "calibration_bins": bins}


def calibrate(logits, labels):
    grid = np.geomspace(0.15, 5.0, 80)
    temperature = min(grid, key=lambda t: scores(logits, labels, float(t))["nll"])
    p = probabilities(logits, float(temperature))
    correct, confidence = p.argmax(axis=1) == labels, p.max(axis=1)
    candidates = []
    for threshold in np.linspace(0.125, 0.99, 90):
        accepted = confidence >= threshold
        count = int(accepted.sum())
        if count >= 32 and float(correct[accepted].mean()) >= 0.95:
            candidates.append((count, -float(threshold), float(threshold)))
    threshold = max(candidates)[2] if candidates else 1.0
    return float(temperature), threshold


def export_model(model, variant, directory, temperature, threshold):
    directory.mkdir(parents=True, exist_ok=False)
    blob, layout = bytearray(), []
    for name, tensor, rows, cols in [
        ("w1", model.first.weight, HIDDEN, FEATURES),
        ("b1", model.first.bias, 1, HIDDEN),
        ("w2", model.last.weight, len(LABELS), HIDDEN),
        ("b2", model.last.bias, 1, len(LABELS)),
    ]:
        array = tensor.detach().cpu().numpy().astype(np.float32).reshape(-1)
        offset, scale = len(blob), 1.0
        if variant != "fp32" and name.startswith("w"):
            scale = max(float(np.abs(array).mean()), 1e-8)
            codes = np.rint(array / scale).clip(-1, 1).astype(np.int8)
            packed = bytearray()
            for start in range(0, len(codes), 5):
                group = [int(x) + 1 for x in codes[start:start + 5]]
                group += [1] * (5 - len(group))
                packed.append(sum(code * (3 ** index) for index, code in enumerate(group)))
            blob.extend(packed)
            encoding = "ternary_base3_5"
        else:
            blob.extend(array.astype("<f4").tobytes())
            encoding = "float32_le"
        layout.append({"name": name, "rows": rows, "cols": cols, "count": len(array),
                       "encoding": encoding, "offset": offset, "bytes": len(blob) - offset, "scale": scale})
    (directory / "weights.bin").write_bytes(blob)
    config = {"schema": "gooo/tiny-ir-decision-model/v1", "variant": variant,
              "feature_dim": FEATURES, "hidden_dim": HIDDEN, "max_bytes": LIMIT,
              "labels": LABELS, "temperature": temperature, "confidence_threshold": threshold,
              "weights_file": "weights.bin", "weights_sha256": sha(blob), "tensors": layout}
    save_json(directory / "model.json", config)
    total_parameters = sum(x["count"] for x in layout)
    return {"weights_bytes": len(blob), "model_json_bytes": (directory / "model.json").stat().st_size,
            "total_parameters": total_parameters, "matrix_parameters": HIDDEN * FEATURES + len(LABELS) * HIDDEN,
            "weight_blob_bits_per_total_parameter": len(blob) * 8 / total_parameters,
            "note": "Packed artifact size is separate from resident execution memory; biases/scales/metadata are not 1.58-bit."}


def exported_logits(directory, x):
    config = json.loads((directory / "model.json").read_text())
    blob = (directory / "weights.bin").read_bytes()
    tensors = {}
    for t in config["tensors"]:
        raw = blob[t["offset"]:t["offset"] + t["bytes"]]
        if t["encoding"] == "float32_le":
            values = np.frombuffer(raw, dtype="<f4").copy()
        else:
            decoded = []
            for code in raw:
                for _ in range(5):
                    decoded.append(code % 3 - 1)
                    code //= 3
            values = np.array(decoded[:t["count"]], dtype=np.float32) * np.float32(t["scale"])
        tensors[t["name"]] = values.reshape(t["rows"], t["cols"])
    h = np.maximum(x @ tensors["w1"].T + tensors["b1"], 0)
    return h @ tensors["w2"].T + tensors["b2"]


def train(x, y, cal_x, cal_y, seed, epochs, batch, device, qat):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    model = Model(qat=qat).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)
    tx, ty = torch.from_numpy(x).to(device), torch.from_numpy(y).to(device)
    cx = torch.from_numpy(cal_x).to(device)
    rng = np.random.default_rng(seed)
    history, best, best_loss = [], None, float("inf")
    start = time.perf_counter()
    max_allocated, max_driver = 0, 0
    for epoch in range(epochs):
        order = rng.permutation(len(y))
        losses = []
        model.train()
        for begin in range(0, len(order), batch):
            index = torch.from_numpy(order[begin:begin + batch]).to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(tx[index])
            loss = nn.functional.cross_entropy(logits, ty[index])
            loss.backward()
            optimizer.step()
            losses.append(float(loss.detach().cpu()))
        model.eval()
        with torch.no_grad():
            cal = model(cx).cpu().numpy()
        value = scores(cal, cal_y, 1.0)["nll"]
        if value < best_loss:
            best_loss = value
            best = {key: tensor.detach().cpu().clone() for key, tensor in model.state_dict().items()}
        if device == "mps":
            max_allocated = max(max_allocated, torch.mps.current_allocated_memory())
            max_driver = max(max_driver, torch.mps.driver_allocated_memory())
        history.append({"epoch": epoch + 1, "training_loss": sum(losses) / len(losses), "calibration_nll": value})
    if device == "mps":
        torch.mps.synchronize()
    duration = time.perf_counter() - start
    model.cpu()
    model.load_state_dict(best)
    return model, {"seed": seed, "epochs": epochs, "batch_size": batch, "steps": epochs * math.ceil(len(y) / batch),
                   "device": device, "training_loop_wall_seconds": duration, "mps_allocated_peak_sampled_bytes": max_allocated,
                   "mps_driver_peak_sampled_bytes": max_driver, "process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
                   "selection": "lowest calibration NLL; no test access during training", "history": history,
                   "timing_note": "Whole loop includes training, calibration forward, CPU transfers and scoring; not pure GPU kernel time",
                   "memory_note": "MPS end-of-epoch samples; process RSS is cumulative across variants, not isolated variant peak"}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--epochs", type=int, default=120)
    parser.add_argument("--batch", type=int, default=256)
    parser.add_argument("--seed", type=int, default=20360930)
    parser.add_argument("--validate-only", action="store_true")
    args = parser.parse_args()
    data_raw, rows, dataset_pins = validate_frozen_dataset(args.dataset)
    require(1 <= args.epochs <= 200 and 1 <= args.batch <= 512, "pilot budget outside bounds")
    if args.validate_only:
        print(json.dumps({"status": "FROZEN_DATASET_VERIFIED", **dataset_pins, "optimization_level": sys.flags.optimize,
                          "model_training_calls": 0}))
        return
    require(torch.backends.mps.is_available(), "This pilot explicitly requires available MPS GPU")
    args.output.mkdir(parents=True, exist_ok=False)
    source = Path(__file__).read_bytes()
    (args.output / "training-source.py").write_bytes(source)
    require(all(row["label"] in LABELS and row["split"] in ("train", "calibration", "test") for row in rows), "invalid split/label")
    split_rows = {split: [r for r in rows if r["split"] == split] for split in ("train", "calibration", "test")}
    prompt_sets = {s: {r["text"] for r in rs} for s, rs in split_rows.items()}
    groups = {s: {r["template_id"] for r in rs} for s, rs in split_rows.items()}
    for a, b in (("train", "calibration"), ("train", "test"), ("calibration", "test")):
        require(not prompt_sets[a] & prompt_sets[b] and not groups[a] & groups[b], "split leakage")
    xs = {s: np.stack([features(r["text"]) for r in rs]) for s, rs in split_rows.items()}
    ys = {s: np.array([LABELS.index(r["label"]) for r in rs], dtype=np.int64) for s, rs in split_rows.items()}
    pre = {"schema": "gooo/tiny-ir-decision-training-preexecution/v1", "training_source_sha256": sha(source),
           **dataset_pins, "python_optimization_level": sys.flags.optimize,
           "dataset_sha256": sha(data_raw), "torch_version": torch.__version__, "numpy_version": np.__version__,
           "seed": args.seed, "epochs": args.epochs, "batch_size": args.batch,
           "hardware": {"chip": "Apple M4", "gpu_cores": 10, "unified_memory_gib": 16},
           "split_counts": {s: len(rs) for s, rs in split_rows.items()}, "data_provenance": "Public synthetic Go-generated templates only",
           "gpu_determinism": "seeded; bit-identical cross-device training is not asserted",
           "variants": ["fp32", "ptq_ternary", "qat_ternary"], "test_used_for_training_or_calibration": False}
    save_json(args.output / "preexecution.json", pre)
    torch.set_num_threads(2)
    models, training = {}, {}
    for variant, qat in (("fp32", False), ("qat_ternary", True)):
        models[variant], training[variant] = train(xs["train"], ys["train"], xs["calibration"], ys["calibration"],
                                                   args.seed, args.epochs, args.batch, "mps", qat)
    models["ptq_ternary"] = models["fp32"]
    results, parity = {}, []
    for variant in ("fp32", "ptq_ternary", "qat_ternary"):
        directory = args.output / "models" / variant
        export_model(models[variant], variant, directory, 1.0, 0.0)
        calibration_logits = exported_logits(directory, xs["calibration"])
        temperature, threshold = calibrate(calibration_logits, ys["calibration"])
        config = json.loads((directory / "model.json").read_text())
        config.update(temperature=temperature, confidence_threshold=threshold)
        save_json(directory / "model.json", config)
        variant_scores = {s: scores(exported_logits(directory, xs[s]), ys[s], temperature) for s in xs}
        tp = probabilities(exported_logits(directory, xs["test"]), temperature)
        accepted = tp.max(axis=1) >= threshold
        correct = tp.argmax(axis=1) == ys["test"]
        matrix_count = HIDDEN * FEATURES + len(LABELS) * HIDDEN
        all_count = matrix_count + HIDDEN + len(LABELS)
        results[variant] = {"scores": variant_scores, "temperature": temperature, "confidence_threshold": threshold,
                            "test_selective": {"planned": len(accepted), "accepted": int(accepted.sum()),
                                               "correct": int((correct & accepted).sum()),
                                               "accuracy": float(correct[accepted].mean()) if accepted.any() else None},
                            "weights_bytes": (directory / "weights.bin").stat().st_size,
                            "model_json_bytes": (directory / "model.json").stat().st_size,
                            "total_parameters": all_count, "matrix_parameters": matrix_count,
                            "weight_blob_bits_per_total_parameter": (directory / "weights.bin").stat().st_size * 8 / all_count,
                            "storage_note": "Matrix codes use 1.6 physical bits/weight; 1.585 is theoretical alphabet entropy. Biases, scales, metadata and runtime memory are separate."}
        logits = exported_logits(directory, xs["test"])
        for index in range(min(32, len(split_rows["test"]))):
            row = split_rows["test"][index]
            parity.append({"variant": variant, "text": row["text"], "features": xs["test"][index].tolist(),
                           "logits": logits[index].tolist(), "probabilities": tp[index].tolist(),
                           "selected_label": LABELS[int(tp[index].argmax())]})
    report = {"schema": "gooo/tiny-ir-decision-training-report/v1", "status": "TRAINED_AND_EXPORTED",
              "preexecution_sha256": sha((args.output / "preexecution.json").read_bytes()), "training": training,
              "variants": results, "limitations": ["Closed eight-operation synthetic task; not arbitrary natural language compilation",
                                                       "One seed and template-group test; no unseen operation family evaluation",
                                                       "QAT is ternary-weight MLP with float activations, not a reproduced BitNet Transformer",
                                                       "Python/exported scores pending independent Go inference parity and compiler-route replay",
                                                       "Sampled memory includes tooling and is not packed-artifact size"]}
    save_json(args.output / "report.json", report)
    save_json(args.output / "go-parity.json", {"schema": "gooo/tiny-ir-decision-parity/v1", "rows": parity})
    print(json.dumps({"status": report["status"], "training_loop_seconds": {k: v["training_loop_wall_seconds"] for k, v in training.items()},
                      "test_accuracy": {k: v["scores"]["test"]["accuracy"] for k, v in results.items()}}))


if __name__ == "__main__":
    main()