"""第 38 章：真实训练、权重 INT8/INT4、QAT、剪枝与蒸馏；CPU 小实验。"""

import argparse
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import platform
from tempfile import TemporaryDirectory
import torch
from torch import nn
from torch.nn import functional as F
from digits_project import make_data

CONFIG = dict(teacher_width=64, student_width=16, epochs=80, qat_epochs=20,
              batch_size=64, lr=.003, qat_lr=.0003, weight_decay=.0001,
              temperature=2., alpha=.7, teacher_seed=2038, student_seed=2039,
              order_seed=2040, max_accuracy_drop=.01, max_state_bytes=6000)


def quantize_rows(weight, bits):
    """对 Linear 的每个输出通道求一个 scale；使用对称范围，保留一个未用码。"""
    if bits not in (4, 8) or weight.ndim != 2 or not torch.isfinite(weight).all():
        raise ValueError("需要有限二维权重，位数为 4 或 8")
    limit = 2**(bits-1)-1
    maximum = weight.detach().abs().amax(dim=1, keepdim=True)
    scale = torch.where(maximum > 0, maximum/limit, torch.ones_like(maximum))
    codes = torch.round(weight.detach()/scale).clamp(-limit, limit).to(torch.int8)
    return codes, scale


def fake_weight(weight, bits):
    codes, scale = quantize_rows(weight, bits)
    rounded = codes.to(weight.dtype)*scale
    # 前向使用舍入后的值；反向对浮点主权重采用恒等直通估计，不对 scale 求导。
    return weight+(rounded-weight).detach()


def pack_int4(codes):
    if codes.dtype != torch.int8 or torch.any(codes < -7) or torch.any(codes > 7):
        raise ValueError("本例的对称 4 位码范围为 -7 到 7")
    unsigned = (codes.flatten().to(torch.int16)+7).to(torch.uint8)
    if unsigned.numel() % 2:
        unsigned = torch.cat([unsigned, torch.zeros(1, dtype=torch.uint8)])
    return unsigned[::2] | (unsigned[1::2] << 4)


def unpack_int4(packed, count):
    if packed.dtype != torch.uint8 or count < 1 or packed.numel() != (count+1)//2:
        raise ValueError("压缩字节数与原元素数不匹配")
    unsigned = torch.stack([packed & 15, packed >> 4], dim=1).flatten()[:count]
    if torch.any(unsigned > 14):
        raise ValueError("遇到本格式未使用的码 15")
    return (unsigned.to(torch.int16)-7).to(torch.int8)


class PackedLinear(nn.Module):
    def __init__(self, layer, bits):
        super().__init__()
        codes, scale = quantize_rows(layer.weight, bits)
        self.bits = bits
        self.rows, self.columns = layer.weight.shape
        self.register_buffer("codes", pack_int4(codes) if bits == 4 else codes)
        self.register_buffer("scale", scale.clone())
        self.register_buffer("bias", layer.bias.detach().clone())

    def weight(self):
        codes = (unpack_int4(self.codes, self.rows*self.columns).reshape(self.rows, self.columns)
                 if self.bits == 4 else self.codes)
        return codes.float()*self.scale

    def forward(self, x):
        # 真正紧凑保存，但每次解包回 FP32 计算；这里没有高效整数矩阵内核。
        return F.linear(x, self.weight(), self.bias)


class Classifier(nn.Module):
    def __init__(self, width, fake_bits=None):
        super().__init__()
        self.width, self.fake_bits = width, fake_bits
        self.first, self.last = nn.Linear(64, width), nn.Linear(width, 10)

    def forward(self, x):
        x = x.flatten(1)
        for index, layer in enumerate([self.first, self.last]):
            x = (F.linear(x, fake_weight(layer.weight, self.fake_bits), layer.bias)
                 if self.fake_bits is not None else layer(x))
            if index == 0:
                x = x.relu()
        return x


def pack_model(model, bits):
    result = deepcopy(model)
    result.fake_bits = None
    result.first, result.last = PackedLinear(model.first, bits), PackedLinear(model.last, bits)
    return result.eval()


def state_bytes(model):
    return sum(t.numel()*t.element_size() for t in model.state_dict().values())


def save_model(model, path):
    if model.fake_bits is not None:
        raise ValueError("请先把 QAT 模型转换为 PackedLinear 再导出")
    bits = model.first.bits if isinstance(model.first, PackedLinear) else None
    torch.save(dict(format=1, width=model.width, bits=bits,
                    state=model.state_dict()), path)


def load_model(path):
    data = torch.load(path, map_location="cpu", weights_only=True)
    if data["format"] != 1 or data["width"] not in (16, 64) or data["bits"] not in (None, 4, 8):
        raise ValueError("不支持的模型格式")
    model = Classifier(data["width"])
    if data["bits"] is not None:
        model = pack_model(model, data["bits"])
    model.load_state_dict(data["state"], strict=True)
    return model.eval()


def distillation_loss(student, teacher, labels, temperature=2., alpha=.7):
    hard = F.cross_entropy(student, labels)
    soft = F.kl_div(F.log_softmax(student/temperature, dim=1),
                   F.softmax(teacher.detach()/temperature, dim=1),
                   reduction="batchmean")
    return (1-alpha)*hard+alpha*temperature**2*soft


@torch.inference_mode()
def evaluate(model, dataset):
    model.eval()
    x, y = dataset
    logits = model(x)
    prediction = logits.argmax(1)
    return dict(loss=F.cross_entropy(logits, y).item(),
                accuracy=(prediction == y).float().mean().item(),
                predictions=prediction.tolist())


def fit(model, train, validation, teacher=None, epochs=None, lr=None):
    epochs = CONFIG["epochs"] if epochs is None else epochs
    lr = CONFIG["lr"] if lr is None else lr
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=CONFIG["weight_decay"])
    generator = torch.Generator().manual_seed(CONFIG["order_seed"])
    x, y = train
    teacher_logits = None
    if teacher is not None:
        teacher.eval()
        with torch.no_grad():
            teacher_logits = teacher(x).detach()
    best_state, best_loss, best_epoch = deepcopy(model.state_dict()), evaluate(model, validation)["loss"], 0
    history = []
    for epoch in range(1, epochs+1):
        model.train()
        for ids in torch.randperm(len(x), generator=generator).split(CONFIG["batch_size"]):
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[ids])
            loss = (F.cross_entropy(logits, y[ids]) if teacher is None else
                    distillation_loss(logits, teacher_logits[ids], y[ids],
                                      CONFIG["temperature"], CONFIG["alpha"]))
            if not torch.isfinite(loss):
                raise RuntimeError("训练损失非有限")
            loss.backward()
            optimizer.step()
        score = evaluate(model, validation)
        history.append(dict(epoch=epoch, loss=score["loss"], accuracy=score["accuracy"]))
        if score["loss"] < best_loss:
            best_state, best_loss, best_epoch = deepcopy(model.state_dict()), score["loss"], epoch
    model.load_state_dict(best_state)
    return model.eval(), best_epoch, history


def prune_weights(model, fraction=.9):
    result = deepcopy(model)
    weights = [result.first.weight, result.last.weight]
    flat = torch.cat([w.detach().flatten() for w in weights])
    mask = torch.ones_like(flat)
    mask[flat.abs().argsort()[:int(fraction*len(flat))]] = 0
    start = 0
    with torch.no_grad():
        for weight in weights:
            count = weight.numel()
            weight.mul_(mask[start:start+count].reshape(weight.shape))
            start += count
    return result


def prune_units(model, train_x, keep=16):
    # 只用训练输入估计隐藏单元贡献，选择后整行/整列删除，不做额外微调。
    with torch.inference_mode():
        activation = model.first(train_x.flatten(1)).relu().mean(0)
        importance = activation*model.last.weight.abs().sum(0)
        chosen = importance.topk(keep).indices.sort().values
    result = Classifier(keep)
    with torch.no_grad():
        result.first.weight.copy_(model.first.weight[chosen])
        result.first.bias.copy_(model.first.bias[chosen])
        result.last.weight.copy_(model.last.weight[:, chosen])
        result.last.bias.copy_(model.last.bias)
    return result.eval()


def run():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    raw, ids = make_data()
    datasets = [(x.float()/16, y) for x, y in raw]
    train, validation, test = datasets
    torch.manual_seed(CONFIG["teacher_seed"])
    teacher, teacher_epoch, teacher_history = fit(Classifier(64), train, validation)
    torch.manual_seed(CONFIG["student_seed"])
    initial = Classifier(16)
    hard, hard_epoch, hard_history = fit(deepcopy(initial), train, validation)
    student, student_epoch, student_history = fit(deepcopy(initial), train, validation, teacher)
    qat = deepcopy(teacher)
    qat.fake_bits = 4
    qat, qat_epoch, qat_history = fit(qat, train, validation, epochs=CONFIG["qat_epochs"],
                                     lr=CONFIG["qat_lr"])
    models = dict(teacher=teacher, ptq8=pack_model(teacher, 8), ptq4=pack_model(teacher, 4),
                  qat4=pack_model(qat, 4), prune90=prune_weights(teacher),
                  units16=prune_units(teacher, train[0]), hard16=hard, distill16=student)
    validation_scores = {name: evaluate(model, validation) for name, model in models.items()}
    eligible = [name for name, model in models.items()
                if name != "teacher" and state_bytes(model) <= CONFIG["max_state_bytes"]
                and validation_scores[name]["accuracy"] >=
                validation_scores["teacher"]["accuracy"]-CONFIG["max_accuracy_drop"]]
    selected = min(eligible, key=lambda name: (state_bytes(models[name]),
                                              validation_scores[name]["loss"])) if eligible else "teacher"
    # 到这里所有方案和选择均已确定；随后用已在旧章公开的测试划分做能力回归。
    reference = torch.tensor(evaluate(teacher, test)["predictions"])
    target = test[1]
    rows = []
    with TemporaryDirectory() as directory:
        for name, model in models.items():
            path = Path(directory)/"model.pt"
            save_model(model, path)
            restored = load_model(path)
            with torch.inference_mode():
                torch.testing.assert_close(model(test[0]), restored(test[0]), atol=0, rtol=0)
            score = evaluate(restored, test)
            prediction = torch.tensor(score["predictions"])
            regressed = (reference == target) & (prediction != target)
            fixed = (reference != target) & (prediction == target)
            weights = [model.first.weight(), model.last.weight()] if isinstance(model.first, PackedLinear) else [
                model.first.weight, model.last.weight]
            rows.append(dict(name=name, state_bytes=state_bytes(model),
                             weight_values=sum(w.numel() for w in weights),
                             nonzero_weights=sum(int((w != 0).sum()) for w in weights),
                             validation=validation_scores[name], test=score,
                             regressions=ids[2][regressed.numpy()].tolist(),
                             fixes=ids[2][fixed.numpy()].tolist(),
                             class_correct=[int(((prediction == target) & (target == digit)).sum())
                                            for digit in range(10)]))
    digest = hashlib.sha256()
    for x, y in datasets:
        digest.update(x.numpy().tobytes())
        digest.update(y.numpy().tobytes())
    report = dict(complete=True, config=CONFIG, selected=selected, data_sha256=digest.hexdigest(),
                  environment=dict(python=platform.python_version(), torch=torch.__version__,
                                   device="cpu", dtype="float32", threads=1),
                  split_ids=[part.tolist() for part in ids], test_targets=target.tolist(),
                  class_counts=[int((target == digit).sum()) for digit in range(10)],
                  epochs=dict(teacher=teacher_epoch, hard16=hard_epoch,
                              distill16=student_epoch, qat4=qat_epoch),
                  history=dict(teacher=teacher_history, hard16=hard_history,
                               distill16=student_history, qat4=qat_history), rows=rows)
    return report, models


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--save-dir", type=Path)
    args = parser.parse_args()
    report, models = run()
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
    if args.save_dir:
        args.save_dir.mkdir(parents=True, exist_ok=True)
        for name, model in models.items():
            save_model(model, args.save_dir/f"{name}.pt")
    print("沿用数字划分：训练/验证/回归测试 = 1077 / 360 / 360")
    print("按验证交叉熵选轮次：", report["epochs"])
    for row in report["rows"]:
        print(f"{row['name']}：状态={row['state_bytes']} 字节；"
              f"验证准确率={row['validation']['accuracy']:.3f}；"
              f"测试准确率={row['test']['accuracy']:.3f}；"
              f"退步/修正={len(row['regressions'])}/{len(row['fixes'])}")
    print("按预设验证条件选择：", report["selected"])
    print("全部模型保存/加载后的分数一致；INT4 每字节实际存两个码。")
    print("量化层解包回 FP32 运算，未测整数内核加速；状态字节不等于峰值内存。")


if __name__ == "__main__":
    main()