"""第 36 章：两个真实 CPU 进程同步梯度，并检查激活重算与数值精度。"""

from copy import deepcopy
from datetime import timedelta
import json
import os
import sys
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel
from torch.utils.checkpoint import checkpoint


def regression_data():
    x = torch.arange(1, 6, dtype=torch.float64).reshape(-1, 1)
    return x, 2*x+1


def zero_model():
    model = nn.Linear(1, 1).double()
    nn.init.zeros_(model.weight)
    nn.init.zeros_(model.bias)
    return model


def full_batch_update():
    x, y = regression_data()
    model = zero_model()
    optimizer = torch.optim.SGD(model.parameters(), lr=.01)
    nn.functional.mse_loss(model(x), y).backward()
    optimizer.step()
    return torch.cat([p.detach().flatten() for p in model.parameters()])


def worker(rank, rendezvous, output_dir):
    # 本例的两个进程都在同一台电脑；默认走回环接口，保留显式网卡设置。
    if sys.platform in ("darwin", "linux"):
        os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo0" if sys.platform == "darwin" else "lo")
    # Gloo 的启动日志留在各进程日志文件；若进程出错，spawn 仍会抛出异常。
    with open(Path(output_dir)/f"rank-{rank}.log", "w") as log:
        os.dup2(log.fileno(), 1)
        os.dup2(log.fileno(), 2)
        torch.set_num_threads(1)
        dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2,
                                timeout=timedelta(seconds=30))
        try:
            x, y = regression_data()
            part = slice(0, 2) if rank == 0 else slice(2, 5)
            local_x, local_y = x[part], y[part]
            total = torch.tensor(float(len(local_x)), dtype=torch.float64)
            dist.all_reduce(total, op=dist.ReduceOp.SUM)
            results = []
            for correct in [True, False]:
                model = DistributedDataParallel(zero_model())
                optimizer = torch.optim.SGD(model.parameters(), lr=.01)
                errors = (model(local_x)-local_y).square()
                # 默认 DDP 平均各 rank 的梯度；因此局部和乘 world_size/全局数量。
                loss = errors.sum()*dist.get_world_size()/total if correct else errors.mean()
                loss.backward()
                optimizer.step()
                values = torch.cat([p.detach().flatten() for p in model.parameters()])
                results.append(values.tolist())
            (Path(output_dir)/f"rank-{rank}.json").write_text(json.dumps(results))
        finally:
            dist.destroy_process_group()


def distributed_demo():
    if not dist.is_available() or not dist.is_gloo_available():
        raise RuntimeError("本实验需要启用 Gloo 的 PyTorch，不以本地求平均代替进程通信")
    with TemporaryDirectory() as directory:
        rendezvous = (Path(directory)/"rendezvous").as_uri()
        mp.spawn(worker, args=(rendezvous, directory), nprocs=2, join=True)
        rows = [json.loads((Path(directory)/f"rank-{rank}.json").read_text())
                for rank in range(2)]
    first, second = [torch.tensor(row, dtype=torch.float64) for row in rows]
    assert torch.equal(first, second), "两个进程更新后参数不同"
    reference = full_batch_update()
    assert torch.allclose(first[0], reference, atol=1e-12, rtol=0)
    assert not torch.allclose(first[1], reference, atol=1e-12, rtol=0)
    return reference, first


class CountedBlock(nn.Module):
    def __init__(self, width=64):
        super().__init__()
        self.layers = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(),
                                    nn.Linear(4*width, width))
        self.calls = 0

    def forward(self, x):
        self.calls += 1
        return x+self.layers(x)


def checkpoint_run(model, x, enabled):
    model = deepcopy(model)
    x = x.clone().requires_grad_()
    parameter_storages = {p.untyped_storage().data_ptr() for p in model.parameters()}
    saved = {}
    def pack(tensor):
        storage = tensor.untyped_storage()
        pointer = storage.data_ptr()
        if storage.nbytes() and pointer not in parameter_storages:
            saved[pointer] = storage.nbytes()
        return tensor
    def unpack(tensor):
        return tensor
    # 只观察前向为反向保存的非参数 storage；不把它叫进程峰值或显存峰值。
    with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
        output = x
        for block in model:
            output = checkpoint(block, output, use_reentrant=False) if enabled else block(output)
        loss = output.square().mean()
    saved_bytes = sum(saved.values())
    loss.backward()
    gradients = torch.cat([p.grad.flatten() for p in model.parameters()])
    return dict(bytes=saved_bytes, calls=sum(b.calls for b in model),
                output=output.detach(), input_gradient=x.grad, gradients=gradients)


def checkpoint_demo():
    torch.manual_seed(2026)
    model = nn.ModuleList([CountedBlock() for _ in range(4)]).double()
    x = torch.randn(16, 64, dtype=torch.float64)
    regular = checkpoint_run(model, x, False)
    recomputed = checkpoint_run(model, x, True)
    for field in ["output", "input_gradient", "gradients"]:
        assert torch.allclose(regular[field], recomputed[field], atol=1e-12, rtol=1e-12)
    assert recomputed["bytes"] < regular["bytes"]
    assert recomputed["calls"] > regular["calls"]
    return regular, recomputed


def tensor_parallel_algebra():
    torch.manual_seed(2027)
    x = torch.randn(3, 4, dtype=torch.float64)
    w = torch.randn(4, 6, dtype=torch.float64)
    column_split = torch.cat([x@part for part in w.split(3, dim=1)], dim=1)
    row_split = sum(a@b for a, b in zip(x.split(2, dim=1), w.split(2, dim=0)))
    return (torch.allclose(x@w, column_split, atol=1e-12),
            torch.allclose(x@w, row_split, atol=1e-12))


def precision_demo():
    value = torch.tensor([1.001, 70000., 1e-8], dtype=torch.float32)
    representations = {name: value.to(dtype).float().tolist() for name, dtype in
                       [("FP32", torch.float32), ("FP16", torch.float16),
                        ("BF16", torch.bfloat16)]}
    torch.manual_seed(2028)
    model = nn.Linear(8, 4)
    x = torch.randn(2, 8)
    with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
        output = model(x)
        loss = output.float().square().mean()
    loss.backward()
    assert all(torch.isfinite(p.grad).all() for p in model.parameters())
    return representations, str(output.dtype), str(model.weight.dtype), str(model.weight.grad.dtype)


def main():
    torch.set_num_threads(1)
    reference, distributed = distributed_demo()
    print("两个真实 CPU/Gloo 进程；本地样本数分别为 2、3。")
    print("整批更新 [w, b]：", [round(v, 6) for v in reference.tolist()])
    print("按全局数量加权：", [round(v, 6) for v in distributed[0].tolist()])
    print("错误地平均本地均值：", [round(v, 6) for v in distributed[1].tolist()])
    print("两个进程参数一致；正确加权与整批更新一致。")
    regular, recomputed = checkpoint_demo()
    print("前向保存的非参数 storage 字节：", regular["bytes"], "→", recomputed["bytes"])
    print("模块前向进入次数：", regular["calls"], "→", recomputed["calls"])
    print("检查点前后输出、输入梯度和参数梯度一致。")
    print("列切分拼接、行切分求和均与原乘积一致：", all(tensor_parallel_algebra()))
    values, output, weight, gradient = precision_demo()
    for name, row in values.items():
        print(name, [format(v, ".6g") for v in row])
    print("CPU autocast 输出 / 参数 / 梯度：", output, weight, gradient)


if __name__ == "__main__":
    main()