"""第 42 章：8×8 数字上的自编码、条件 DDPM 与条件流匹配；CPU 从零训练。"""

from argparse import ArgumentParser
from copy import deepcopy
from hashlib import sha256
import json
from math import pi
from pathlib import Path
import platform
from time import perf_counter

import numpy as np
from sklearn.linear_model import LogisticRegression
import torch
from torch import nn

from digits_project import make_data


STEPS = 64
EPOCHS = {"autoencoder": 80, "ddpm": 120, "flow": 120}


def tensor_hash(tensor):
    return sha256(tensor.detach().cpu().contiguous().numpy().tobytes()).hexdigest()


def schedule(steps=STEPS):
    # 余弦 alpha_bar 曲线；限制最后一步 beta，避免除以零。
    grid = torch.arange(steps + 1, dtype=torch.float64) / steps
    curve = torch.cos((grid + 0.008) / 1.008 * pi / 2).square()
    curve = curve / curve[0]
    beta = (1 - curve[1:] / curve[:-1]).clamp(max=0.999).float()
    alpha_bar = torch.cat([torch.ones(1), (1 - beta).cumprod(0)])
    return beta, alpha_bar


def add_noise(clean, noise, level, alpha_bar):
    a = alpha_bar[level, None]
    return a.sqrt() * clean + (1 - a).sqrt() * noise


def flow_pair(clean, noise, time):
    return (1 - time[:, None]) * noise + time[:, None] * clean, clean - noise


class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(64, 32), nn.SiLU(), nn.Linear(32, 8))
        self.decoder = nn.Sequential(nn.Linear(8, 32), nn.SiLU(),
                                     nn.Linear(32, 64), nn.Sigmoid())

    def forward(self, images):
        return self.decoder(self.encoder(images))


class Field(nn.Module):
    def __init__(self):
        super().__init__()
        self.labels = nn.Embedding(10, 16)
        self.network = nn.Sequential(nn.Linear(64 + 9 + 16, 128), nn.SiLU(),
                                     nn.Linear(128, 128), nn.SiLU(),
                                     nn.Linear(128, 64))

    def forward(self, images, time, labels):
        angles = time[:, None] * (2 * pi * 2 ** torch.arange(4))
        clock = torch.cat([time[:, None], angles.sin(), angles.cos()], dim=1)
        return self.network(torch.cat([images, clock, self.labels(labels)], dim=1))


def training_pair(kind, clean, generator, alpha_bar):
    noise = torch.randn(clean.shape, generator=generator)
    if kind == "ddpm":
        level = torch.randint(1, STEPS + 1, (len(clean),), generator=generator)
        return add_noise(clean, noise, level, alpha_bar), level / STEPS, noise
    time = torch.rand(len(clean), generator=generator)
    mixed, velocity = flow_pair(clean, noise, time)
    return mixed, time, velocity


def train(kind, training, validation):
    torch.manual_seed(2026 if kind == "autoencoder" else 2030)
    model = Autoencoder() if kind == "autoencoder" else Field()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.002)
    generator = torch.Generator().manual_seed(3031)
    _, alpha_bar = schedule()
    x, y = training
    vx, vy = validation
    if kind != "autoencoder":
        x, vx = 2 * x - 1, 2 * vx - 1
        validation_pair = training_pair(
            kind, vx, torch.Generator().manual_seed(3030), alpha_bar)
    best, history = float("inf"), []
    for epoch in range(1, EPOCHS[kind] + 1):
        model.train()
        total = 0.0
        for ids in torch.randperm(len(x), generator=generator).split(128):
            optimizer.zero_grad(set_to_none=True)
            if kind == "autoencoder":
                output, target = model(x[ids]), x[ids]
            else:
                mixed, time, target = training_pair(kind, x[ids], generator, alpha_bar)
                output = model(mixed, time, y[ids])
            loss = (output - target).square().mean()
            if not torch.isfinite(loss):
                raise RuntimeError("训练出现非有限损失")
            loss.backward()
            optimizer.step()
            total += loss.item() * len(ids)
        model.eval()
        with torch.no_grad():
            if kind == "autoencoder":
                loss = (model(vx) - vx).square().mean().item()
            else:
                mixed, time, target = validation_pair
                loss = (model(mixed, time, vy) - target).square().mean().item()
        history.append([total / len(x), loss])
        if loss < best:
            best, best_epoch, state = loss, epoch, deepcopy(model.state_dict())
    model.load_state_dict(state)
    model.eval()
    fingerprint = sha256(b"".join(
        value.numpy().tobytes() for value in state.values())).hexdigest()
    return model, {"epoch": best_epoch, "validation_mse": best,
                   "history": history, "parameters": sum(p.numel() for p in model.parameters()),
                   "weights_sha256": fingerprint}


def reverse_step(noisy, epsilon, step, beta, alpha_bar, noise):
    b, current, previous = beta[step - 1], alpha_bar[step], alpha_bar[step - 1]
    clean = ((noisy - (1 - current).sqrt() * epsilon) / current.sqrt()).clamp(-1, 1)
    mean = (previous.sqrt() * b * clean
            + (1 - b).sqrt() * (1 - previous) * noisy) / (1 - current)
    variance = b * (1 - previous) / (1 - current)
    return mean + variance.sqrt() * noise


@torch.no_grad()
def sample(kind, model, labels, steps=32, seed=4040):
    generator = torch.Generator().manual_seed(seed)
    x = torch.randn((len(labels), 64), generator=generator)
    initial_hash = tensor_hash(x)
    snapshots = [x[0].tolist()]
    positions = [0]
    if kind == "ddpm":
        beta, alpha_bar = schedule()
        for step in range(STEPS, 0, -1):
            time = torch.full((len(x),), step / STEPS)
            epsilon = model(x, time, labels)
            noise = torch.randn(x.shape, generator=generator) if step > 1 else torch.zeros_like(x)
            x = reverse_step(x, epsilon, step, beta, alpha_bar, noise)
            if step in [49, 33, 17, 1]:
                snapshots.append(x[0].tolist())
                positions.append(STEPS - step + 1)
    else:
        for step in range(steps):
            time = torch.full((len(x),), step / steps)
            x = x + model(x, time, labels) / steps
            if (step + 1) % (steps // 4) == 0:
                snapshots.append(x[0].tolist())
                positions.append(step + 1)
    return ((x + 1) / 2).clamp(0, 1), {
        "initial_noise_sha256": initial_hash, "snapshots": snapshots,
        "snapshot_steps": positions, "evaluations": STEPS if kind == "ddpm" else steps,
        "raw_outside_fraction": ((x < -1) | (x > 1)).float().mean().item()}


def run():
    started = perf_counter()
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    raw, ids = make_data()
    data = [(x.flatten(1).float() / 16, y) for x, y in raw]
    training, validation, test = data
    report = {"split_ids": [index.tolist() for index in ids],
              "environment": {"python": platform.python_version(),
                              "torch": torch.__version__, "numpy": np.__version__,
                              "platform": platform.platform(), "threads": 1},
              "config": {"epochs": EPOCHS, "batch_size": 128, "lr": 0.002,
                         "ddpm_steps": STEPS, "flow_steps": 32,
                         "sample_seed": 4040, "train_seed": 3031,
                         "validation_seed": 3030},
              "source_sha256": sha256(Path(__file__).read_bytes()).hexdigest(),
              "data_sha256": [tensor_hash(x) for x, _ in data]}
    ae, record = train("autoencoder", training, validation)
    with torch.no_grad():
        reconstruction = ae(test[0])
        record["test_mse"] = (reconstruction - test[0]).square().mean().item()
        record["mean_baseline_mse"] = (
            training[0].mean(0) - test[0]).square().mean().item()
        chosen = [int(torch.where(test[1] == label)[0][0]) for label in range(10)]
        record["image_ids"] = ids[2][chosen].tolist()
        record["originals"] = test[0][chosen].tolist()
        record["reconstructions"] = reconstruction[chosen].tolist()
        random_latent = torch.randn((10, 8), generator=torch.Generator().manual_seed(4041))
        record["random_latent_decodes"] = ae.decoder(random_latent).tolist()
    report["autoencoder"] = record
    classifier = LogisticRegression(C=1, max_iter=1000).fit(*[v.numpy() for v in training])
    report["classifier_real_accuracy"] = float(classifier.score(
        test[0].numpy(), test[1].numpy()))
    labels = torch.arange(10).repeat_interleave(10)
    report["sample_labels"] = labels.tolist()
    for kind in ["ddpm", "flow"]:
        model, record = train(kind, training, validation)
        samples, trajectory = sample(kind, model, labels)
        predicted = classifier.predict(samples.numpy())
        record.update(trajectory)
        record.update({"samples": samples.tolist(), "predicted": predicted.tolist(),
                       "condition_correct": int((predicted == labels.numpy()).sum())})
        report[kind] = record
    _, alpha_bar = schedule()
    clean = 2 * test[0][chosen[3]:chosen[3] + 1] - 1
    noise = torch.randn(clean.shape, generator=torch.Generator().manual_seed(4042))
    report["forward"] = {"image_id": int(ids[2][chosen[3]]), "levels": [0, 16, 32, 48, 64],
                         "images": [add_noise(clean, noise, torch.tensor([i]), alpha_bar)[0].tolist()
                                    for i in [0, 16, 32, 48, 64]],
                         "alpha_bar": alpha_bar.tolist()}
    report["elapsed_seconds"] = perf_counter() - started
    return report


def show(report):
    ae = report["autoencoder"]
    print("训练/验证/测试：1077 / 360 / 360；图像：8×8")
    print(f"自编码器：最佳轮次={ae['epoch']}，测试 MSE={ae['test_mse']:.4f}")
    print(f"训练均值图基线 MSE={ae['mean_baseline_mse']:.4f}")
    print(f"代理分类器真实测试准确率={report['classifier_real_accuracy']:.3f}")
    for kind in ["ddpm", "flow"]:
        row = report[kind]
        print(f"{kind}：最佳轮次={row['epoch']}，验证目标 MSE={row['validation_mse']:.4f}，"
              f"条件识别={row['condition_correct']}/100，采样调用={row['evaluations']}")
    print("条件识别只检查生成数字是否被分类器认作指定类别，不代表整体生成质量。")


if __name__ == "__main__":
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="保存本次完整训练、像素和评价记录")
    args = parser.parse_args()
    result = run()
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
    show(result)