"""Model Refactoring Handbook: a reproducible CPU learning lab.

Original example code. Tested with Python 3.13.5 and PyTorch 2.10.0+cpu.
This is a synthetic four-class classifier, not an audio or mobile benchmark.
Run: python lab.py --out ../evidence/run
"""
from __future__ import annotations

import argparse
import copy
import hashlib
import json
import math
import platform
import random
import statistics
import sys
import time
from pathlib import Path
from typing import Any

import torch
from torch import Tensor, nn
from torch.nn import functional as F

SEED = 20260924


class ResidualFFN(nn.Module):
    """Shape-preserving residual block with a pruneable intermediate width."""
    def __init__(self, width: int = 32, expansion: int = 96) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(width)
        self.up = nn.Linear(width, expansion)
        self.down = nn.Linear(expansion, width)

    def forward(self, x: Tensor) -> Tensor:
        return x + 0.25 * self.down(F.gelu(self.up(self.norm(x))))


class RefactorNet(nn.Module):
    """Version 1.0: stem, six residual MLP blocks, and a classifier head."""
    def __init__(self, expansions: tuple[int, ...] = (96,) * 6) -> None:
        super().__init__()
        self.stem = nn.Linear(16, 32)
        self.blocks = nn.ModuleList([ResidualFFN(32, h) for h in expansions])
        self.norm = nn.LayerNorm(32)
        self.head = nn.Linear(32, 4)

    def forward(self, x: Tensor) -> Tensor:
        x = F.gelu(self.stem(x))
        for block in self.blocks:
            x = block(x)
        return self.head(self.norm(x))

    def config(self) -> dict[str, Any]:
        return {"architecture": "RefactorNet", "version": "1.0",
                "expansions": [b.up.out_features for b in self.blocks]}


def seed_all(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.use_deterministic_algorithms(True)


def make_data(n: int, seed: int) -> tuple[Tensor, Tensor]:
    """A fixed nonlinear task with disjoint, independently generated splits."""
    x = torch.randn(n, 16, generator=torch.Generator().manual_seed(seed))
    scores = torch.stack([
        1.2*x[:, 0] + 0.6*x[:, 4]*x[:, 5] + torch.sin(x[:, 8]),
        1.2*x[:, 1] + 0.6*x[:, 6]*x[:, 7] + torch.sin(x[:, 9]),
        1.2*x[:, 2] - 0.6*x[:, 4]*x[:, 7] + torch.sin(x[:, 10]),
        1.2*x[:, 3] - 0.6*x[:, 5]*x[:, 6] + torch.sin(x[:, 11]),
    ], dim=1)
    return x, scores.argmax(dim=1)


@torch.inference_mode()
def evaluate(model: nn.Module, data: tuple[Tensor, Tensor]) -> dict[str, float]:
    model.eval()
    x, y = data
    z = model(x)
    if not torch.isfinite(z).all():
        raise ValueError("Non-finite logits")
    return {"loss": F.cross_entropy(z, y).item(),
            "accuracy": (z.argmax(1) == y).float().mean().item()}


def fit(model: nn.Module, train: tuple[Tensor, Tensor], *, epochs: int,
        seed: int, teacher: nn.Module | None = None) -> list[float]:
    """Fresh optimizer after surgery; fixed recovery budget and data order.

    The teacher is read only. The test and validation splits are not passed here.
    """
    seed_all(seed)
    model.train()
    if teacher is not None:
        teacher.eval()
        for p in teacher.parameters():
            p.requires_grad_(False)
    opt = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.0001)
    x, y = train
    rng = torch.Generator().manual_seed(seed)
    history: list[float] = []
    for _ in range(epochs):
        perm = torch.randperm(len(x), generator=rng)
        total = 0.0
        for ids in perm.split(128):
            xb, yb = x[ids], y[ids]
            opt.zero_grad(set_to_none=True)
            logits = model(xb)
            loss = F.cross_entropy(logits, yb)
            if teacher is not None:
                with torch.no_grad():
                    targets = teacher(xb)
                temperature = 2.0
                kd = F.kl_div(
                    F.log_softmax(logits / temperature, dim=-1),
                    F.softmax(targets / temperature, dim=-1),
                    reduction="batchmean",
                ) * temperature**2
                loss = 0.5 * loss + 0.5 * kd
            loss.backward()
            opt.step()
            total += loss.item() * len(ids)
        history.append(total / len(x))
    model.eval()
    return history


def remove_block(model: RefactorNet, index: int) -> RefactorNet:
    if not 0 <= index < len(model.blocks):
        raise IndexError(index)
    student = copy.deepcopy(model)
    student.blocks = nn.ModuleList(
        [b for i, b in enumerate(student.blocks) if i != index]
    )
    return student


@torch.no_grad()
def shrink_ffn(block: ResidualFFN, keep: Tensor) -> ResidualFFN:
    """Delete corresponding rows of up and columns of down.

    Preconditions: ordinary dense Linear -> elementwise GELU -> Linear;
    no gated branch, tied parameters, or normalization over the hidden axis.
    """
    if keep.ndim != 1 or keep.dtype != torch.long or keep.numel() == 0:
        raise ValueError("keep must be a nonempty 1D int64 tensor")
    if keep.unique().numel() != keep.numel():
        raise ValueError("Duplicate hidden indices")
    if keep.min().item() < 0 or keep.max().item() >= block.up.out_features:
        raise IndexError("Hidden index out of range")
    keep = keep.to(block.up.weight.device)
    new = ResidualFFN(block.up.in_features, len(keep)).to(
        device=block.up.weight.device, dtype=block.up.weight.dtype)
    new.norm.load_state_dict(block.norm.state_dict())
    new.up.weight.copy_(block.up.weight[keep])
    new.up.bias.copy_(block.up.bias[keep])
    new.down.weight.copy_(block.down.weight[:, keep])
    new.down.bias.copy_(block.down.bias)
    new.train(block.training)
    return new


@torch.inference_mode()
def width_pruned(model: RefactorNet, calibration: Tensor,
                 keep_count: int = 64) -> tuple[RefactorNet, list[list[int]]]:
    """Activation-weight score is a heuristic, not a proof of redundancy.

    Score channels at the ORIGINAL model's inputs using training-only examples.
    """
    if not 1 <= keep_count <= 96:
        raise ValueError("keep_count must be between 1 and 96")
    model.eval()
    student = copy.deepcopy(model)
    h = F.gelu(model.stem(calibration))
    retained: list[list[int]] = []
    for i, block in enumerate(model.blocks):
        a = F.gelu(block.up(block.norm(h)))
        score = a.abs().mean(0) * block.down.weight.norm(dim=0)
        keep = torch.argsort(score, descending=True, stable=True)[:keep_count]
        keep = keep.sort().values
        retained.append(keep.tolist())
        student.blocks[i] = shrink_ffn(block, keep)
        h = block(h)
    return student, retained


@torch.no_grad()
def factor_linear(layer: nn.Linear, rank: int) -> nn.Sequential:
    """SVD approximation; full rank is only a numerical equivalence check."""
    m, n = layer.weight.shape
    if not 1 <= rank <= min(m, n):
        raise ValueError("rank outside valid range")
    u, s, vh = torch.linalg.svd(layer.weight.float(), full_matrices=False)
    a = nn.Linear(n, rank, bias=False).to(layer.weight)
    b = nn.Linear(rank, m, bias=layer.bias is not None).to(layer.weight)
    a.weight.copy_(vh[:rank].to(layer.weight))
    b.weight.copy_((u[:, :rank] * s[:rank]).to(layer.weight))
    if layer.bias is not None:
        b.bias.copy_(layer.bias)
    return nn.Sequential(a, b).train(layer.training)


def self_tests() -> dict[str, str]:
    seed_all(SEED)
    block = ResidualFFN().eval()
    x = torch.randn(7, 32)
    keep = torch.arange(0, 96, 2)
    pruned = shrink_ffn(block, keep)
    with torch.inference_mode():
        a = F.gelu(block.up(block.norm(x)))
        mask = torch.zeros_like(a)
        mask[:, keep] = 1
        masked = x + 0.25 * block.down(a * mask)
        torch.testing.assert_close(pruned(x), masked, atol=1e-6, rtol=1e-5)
        full = factor_linear(block.up, 32)
        torch.testing.assert_close(full(x), block.up(x), atol=3e-6, rtol=3e-5)
        small = factor_linear(block.up, 12)
        assert small(x).shape == block.up(x).shape
    assert sum(p.numel() for p in pruned.parameters()) < sum(
        p.numel() for p in block.parameters())
    try:
        shrink_ffn(block, torch.tensor([0, 0]))
        raise AssertionError("Duplicate indices accepted")
    except ValueError:
        pass
    # Masks produce zero values, but they do not shrink parameter dimensions.
    from torch.nn.utils import prune
    lin = nn.Linear(16, 32)
    shape = tuple(lin.weight.shape)
    prune.l1_unstructured(lin, "weight", amount=0.5)
    prune.remove(lin, "weight")
    assert tuple(lin.weight.shape) == shape
    assert (lin.weight == 0).sum().item() == 256
    # Layer removal can execute without implying output equivalence.
    assert remove_block(RefactorNet(), 2)(torch.randn(2, 16)).shape == (2, 4)
    return {"masked_vs_structural_ffn": "passed",
            "full_rank_svd_equivalence": "passed",
            "low_rank_svd_shape": "passed",
            "physical_parameter_reduction": "passed",
            "invalid_indices_rejected": "passed",
            "mask_does_not_shrink_tensor": "passed",
            "block_removal_executes": "passed"}


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def benchmark(models: dict[str, nn.Module], x: Tensor) -> dict[str, Any]:
    """Warm, synchronous CPU timings, no profiler and no model loading.

    Same single input, batch one. All models remain resident. Do NOT use
    this experiment to infer process peak memory, Android latency, or thermals.
    """
    results: dict[str, list[float]] = {k: [] for k in models}
    rounds: dict[str, list[float]] = {k: [] for k in models}
    order_log: list[list[str]] = []
    rng = random.Random(SEED)
    with torch.inference_mode():
        for model in models.values():
            model.eval()
            for _ in range(30):
                model(x)
        for _ in range(5):
            order = list(models)
            rng.shuffle(order)
            order_log.append(order)
            for name in order:
                model = models[name]
                samples = []
                for _ in range(200):
                    start = time.perf_counter_ns()
                    model(x)
                    samples.append((time.perf_counter_ns() - start) / 1e6)
                results[name].extend(samples)
                rounds[name].append(statistics.median(samples))
    return {"order": order_log, "warmup_per_model": 30,
            "rounds": 5, "samples_per_round": 200,
            "results": {name: {
                "p50_ms": statistics.median(v),
                "p95_ms": sorted(v)[math.ceil(0.95*len(v))-1],
                "round_medians_ms": rounds[name],
                "raw_ms": v,
            } for name, v in results.items()}}


def run(out: Path) -> None:
    out.mkdir(parents=True, exist_ok=False)
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
    tests = self_tests()
    seed_all(SEED)
    train = make_data(4096, SEED + 1)
    valid = make_data(1024, SEED + 2)
    # The test split is generated only after every candidate has been defined.
    teacher = RefactorNet()
    histories = {"teacher": fit(teacher, train, epochs=24, seed=SEED + 10)}
    baseline = evaluate(teacher, valid)
    ablations = []
    for i in range(len(teacher.blocks)):
        quality = evaluate(remove_block(teacher, i), valid)
        ablations.append({"removed_block": i, **quality,
                          "loss_delta": quality["loss"] - baseline["loss"]})
    selected = min(ablations, key=lambda d: (d["loss_delta"], d["removed_block"]))
    drop = int(selected["removed_block"])
    depth = remove_block(teacher, drop)
    width, keep = width_pruned(teacher, train[0][:256], keep_count=64)
    combined = remove_block(width, drop)
    models: dict[str, RefactorNet] = {
        "baseline": teacher,
        "depth_raw": depth,
        "depth_ce": copy.deepcopy(depth),
        "depth_kd": copy.deepcopy(depth),
        "width_raw": width,
        "width_kd": copy.deepcopy(width),
        "combined_raw": combined,
        "combined_kd": copy.deepcopy(combined),
        "baseline_extra_ce": copy.deepcopy(teacher),
    }
    for name in ("depth_ce", "depth_kd", "width_kd", "combined_kd", "baseline_extra_ce"):
        histories[name] = fit(models[name], train, epochs=8, seed=SEED + 20,
                              teacher=teacher if name.endswith("_kd") else None)
    test = make_data(1024, SEED + 3)
    records: dict[str, Any] = {}
    for name, model in models.items():
        path = out / f"{name}.pt"
        torch.save({"config": model.config(), "state_dict": model.state_dict()}, path)
        reloaded = load_model(path)
        with torch.inference_mode():
            torch.testing.assert_close(model(test[0][:8]), reloaded(test[0][:8]),
                                       atol=0, rtol=0)
        records[name] = {"config": model.config(),
                         "parameters": sum(p.numel() for p in model.parameters()),
                         "parameter_bytes": sum(p.numel()*p.element_size() for p in model.parameters()),
                         "checkpoint_bytes": path.stat().st_size,
                         "checkpoint_sha256": sha256(path),
                         "validation": evaluate(model, valid),
                         "test": evaluate(model, test)}
    tests["all_checkpoint_roundtrips"] = "passed"
    timing = benchmark(models, test[0][:1])
    # A metadata-only hook pass. Output sizes are not live peak activations.
    output_shapes: list[dict[str, Any]] = []
    handles = []
    def record(name: str):
        def hook(module: nn.Module, args: tuple[Tensor, ...], output: Tensor) -> None:
            output_shapes.append({"module": name, "shape": list(output.shape),
                                  "dtype": str(output.dtype),
                                  "output_bytes": output.numel()*output.element_size()})
        return hook
    for name, module in teacher.named_modules():
        if name and not list(module.children()):
            handles.append(module.register_forward_hook(record(name)))
    try:
        with torch.inference_mode():
            teacher(test[0][:1])
    finally:
        for h in handles:
            h.remove()
    cpu = "unknown"
    if Path("/proc/cpuinfo").exists():
        cpu = next((line.split(":", 1)[1].strip() for line in Path("/proc/cpuinfo").read_text().splitlines()
                    if line.startswith("model name")), cpu)
    report = {"experiment": "RefactorNet 1.0 CPU learning lab", "seed": SEED,
              "environment": {"python": sys.version, "torch": torch.__version__,
                              "platform": platform.platform(), "cpu": cpu,
                              "torch_threads": torch.get_num_threads(),
                              "torch_interop_threads": torch.get_num_interop_threads(),
                              "source_sha256": sha256(Path(__file__))},
              "data": {"train": 4096, "validation": 1024, "test": 1024,
                       "calibration": "First 256 training inputs; not held-out test inputs"},
              "protocol": {"teacher_epochs": 24, "recovery_epochs": 8,
                           "batch": 128, "lr": 0.001, "weight_decay": 0.0001,
                           "kd_temperature": 2.0, "kd_weight": 0.5,
                           "width_keep": 64, "depth_selection": "Minimum validation loss delta",
                           "test_used_for_selection": False},
              "tests": tests, "ablations": ablations, "selected_block": drop,
              "width_keep_indices": keep, "models": records, "histories": histories,
              "limitations": ["One seed; synthetic task; not a production quality benchmark",
                              "No Android, NPU, audio, or Urdu evaluation",
                              "No peak RSS or peak activation measurement",
                              "ONNX export and quantization not executed in this environment"]}
    (out / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    (out / "timing.json").write_text(json.dumps(timing, indent=2) + "\n")
    (out / "module_outputs.json").write_text(json.dumps(output_shapes, indent=2) + "\n")
    print(json.dumps({"selected_block": drop, "tests": tests,
          "models": {k: {"parameters": v["parameters"], "test_accuracy": v["test"]["accuracy"],
                         "p50_ms": timing["results"][k]["p50_ms"]} for k, v in records.items()}}, indent=2))


def load_model(path: Path) -> RefactorNet:
    payload = torch.load(path, map_location="cpu", weights_only=True)
    cfg = payload["config"]
    if cfg.get("architecture") != "RefactorNet" or cfg.get("version") != "1.0":
        raise ValueError("Unsupported model configuration")
    model = RefactorNet(tuple(cfg["expansions"]))
    model.load_state_dict(payload["state_dict"], strict=True)
    return model.eval()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, default=Path("run"))
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(json.dumps(self_tests(), indent=2))
    else:
        run(args.out)