"""Reload a checkpoint written by train-mnist.py, inspect it, and evaluate it without training.

Purpose: prove that the .pt file is the model. List what the file holds and how many bytes
         each tensor is, rebuild the network from its class, load the weights, evaluate the
         test set once, and then show what a shape mismatch looks like.
Platform: all (runs on the CPU on purpose: a checkpoint is device-independent)
Minimum memory: 8 GB
Assumes: torch and torchvision are installed; a checkpoint from train-mnist.py exists;
         MNIST is already under ./data (train-mnist.py downloaded it).

Usage: python reload-mnist.py [--checkpoint mnist_mlp.pt] [--hidden 256] [--data ./data]
"""
import argparse
import sys
import zipfile
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class TwoLayerNet(nn.Module):
    """The same architecture as train-mnist.py: the state_dict keys must match it exactly."""

    def __init__(self, hidden: int = 256) -> None:
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(nn.Linear(28 * 28, hidden), nn.ReLU(), nn.Linear(hidden, 10))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(self.flatten(x))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--checkpoint", default="mnist_mlp.pt")
    parser.add_argument("--hidden", type=int, default=256, help="must match the run that wrote the checkpoint")
    parser.add_argument("--data", default="./data")
    args = parser.parse_args()

    path = Path(args.checkpoint)
    if not path.exists():
        sys.exit(f"reload-mnist: {path} does not exist; run train-mnist.py first")

    # 1. The file is a zip archive (torch.save's format since PyTorch 1.6): names and sizes.
    print(f"file: {path}  {path.stat().st_size:,} bytes on disk")
    with zipfile.ZipFile(path) as zf:
        for info in zf.infolist():
            print(f"  {info.filename:40s} {info.file_size:>9,} bytes")

    # 2. The state_dict: an ordered mapping from parameter name to tensor.
    sd = torch.load(path, map_location="cpu")  # weights_only=True is the default: tensors only, no code
    print(f"\nstate_dict: {type(sd).__name__} with {len(sd)} entries")
    total = 0
    for name, t in sd.items():
        nbytes = t.numel() * t.element_size()
        total += nbytes
        print(f"  {name:20s} {str(tuple(t.shape)):12s} {str(t.dtype):14s} {t.numel():>8,} x {t.element_size()} B = {nbytes:>9,} bytes")
    print(f"  {'total':20s} {'':12s} {'':14s} {sum(t.numel() for t in sd.values()):>8,} params   {total:>9,} bytes")

    # 3. Rebuild the network and load the weights into it.
    model = TwoLayerNet(args.hidden)
    try:
        result = model.load_state_dict(sd)  # strict=True: every key and shape must match
    except RuntimeError as err:
        sys.exit("reload-mnist: the checkpoint was written with a different --hidden; "
                 + str(err).splitlines()[1].strip())
    model.eval()
    print(f"\nload_state_dict: missing {list(result.missing_keys)}  unexpected {list(result.unexpected_keys)}")

    # 4. Evaluate the test set once, exactly as train-mnist.py did, without a single training step.
    test_set = datasets.MNIST(args.data, train=False, download=False, transform=transforms.ToTensor())
    loader = DataLoader(test_set, batch_size=1000)
    loss_fn = nn.CrossEntropyLoss()
    total_loss, correct, seen = 0.0, 0, 0
    with torch.no_grad():
        for images, labels in loader:
            logits = model(images)
            total_loss += loss_fn(logits, labels).item() * labels.size(0)
            correct += (logits.argmax(dim=1) == labels).sum().item()
            seen += labels.size(0)
    print(f"test loss {total_loss / seen:.4f}  test acc {correct / seen:.4f}  ({seen:,} images, cpu, no training)")

    # 5. What a mismatch looks like: the same file into a network of a different width.
    wrong_hidden = 128 if args.hidden != 128 else 64
    print(f"\nloading the same file into TwoLayerNet(hidden={wrong_hidden}) ...")
    try:
        TwoLayerNet(wrong_hidden).load_state_dict(sd)
    except RuntimeError as err:
        first_lines = str(err).splitlines()[:2]
        print("RuntimeError:", first_lines[0])
        if len(first_lines) > 1:
            print("   ", first_lines[1].strip()[:110], "...")


if __name__ == "__main__":
    main()