# examples/vit_export.py
"""Reproduce the two saved ViT predictions, crops, and torchvision parity checks.

Keep this file beside vit.py. Run --output-dir vit-demo; it downloads the two
CC0 source photographs and the explicitly named pretrained checkpoint.
"""

import argparse
import hashlib
import json
from pathlib import Path
from urllib.request import urlopen

from PIL import Image
import torch
import torchvision
from torchvision.models import vit_b_16

import vit

PHOTOS = [
    (
        "cat",
        "Chelsea the cat",
        "chelsea.png",
        "Stefan van der Walt",
        "596aa1e7cb875eb79f437e310381d26b338a81c2da23439704a73c4651e8c4bb",
    ),
    (
        "coffee",
        "Cup of coffee",
        "coffee.png",
        "Rachel Michetti",
        "cc02f8ca188b167c775a7101b5d767d1e71792cf762c33d6fa15a4599b5a8de7",
    ),
]
SOURCE_BASE = "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.25.2/skimage/data/"


def digest(content):
    """Hash bytes to identify the inputs used for this run."""
    return hashlib.sha256(content).hexdigest()


def export(output_dir):
    """Export actual outputs and compare each logit with the reference model."""
    output_dir.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(1)
    model, weights = vit.load_pretrained()
    reference = vit_b_16(weights=weights).cpu().float().eval()
    transform = weights.transforms()
    results = []
    for photo_id, title, filename, author, expected_hash in PHOTOS:
        with urlopen(SOURCE_BASE + filename, timeout=60) as response:
            content = response.read()
        if digest(content) != expected_hash:
            raise ValueError("The source photo bytes changed")
        source_path = output_dir / f"source-{photo_id}.png"
        source_path.write_bytes(content)
        with Image.open(source_path) as image:
            crop, batch, logits, top5 = vit.predict(model, weights, image)
            original_shape = [image.height, image.width]
            torch.testing.assert_close(
                batch[0], transform(image.convert("RGB")), rtol=0, atol=0
            )
        with torch.inference_mode():
            reference_logits = reference(batch)
        torch.testing.assert_close(
            logits, reference_logits, rtol=2e-4, atol=2e-5
        )
        input_path = output_dir / f"input-{photo_id}.png"
        crop.save(input_path)
        results.append(
            {
                "id": photo_id,
                "title": title,
                "originalImage": source_path.name,
                "originalShape": original_shape,
                "inputImage": input_path.name,
                "inputShape": list(batch.shape),
                "inputSha256": digest(input_path.read_bytes()),
                "source": {
                    "url": SOURCE_BASE + filename,
                    "sha256": expected_hash,
                    "author": author,
                    "license": "CC0",
                    "licenseUrl": "https://github.com/scikit-image/scikit-image/blob/v0.25.2/skimage/data/_fetchers.py",
                },
                "top5": top5,
                "referenceMaxAbsError": float(
                    (logits - reference_logits).abs().max()
                ),
            }
        )
    checkpoint = (
        Path(torch.hub.get_dir())
        / "checkpoints"
        / weights.url.rsplit("/", 1)[1]
    )
    data = {
        "provenance": {
            "torch": torch.__version__,
            "torchvision": torchvision.__version__,
            "programSha256": digest(Path(vit.__file__).read_bytes()),
            "exportSha256": digest(Path(__file__).read_bytes()),
            "weights": str(weights),
            "weightsUrl": weights.url,
            "weightsSha256": digest(checkpoint.read_bytes()),
            "training": "ImageNet-1k from scratch with a modified DeiT recipe",
            "reference": "torchvision.models.vit_b_16",
            "parameterCount": sum(p.numel() for p in model.parameters()),
        },
        "preprocessing": {
            "resizeSize": transform.resize_size,
            "cropSize": transform.crop_size,
            "interpolation": transform.interpolation.value,
            "antialias": transform.antialias,
            "mean": transform.mean,
            "std": transform.std,
        },
        "photos": results,
    }
    (output_dir / "predictions.json").write_text(
        json.dumps(data, indent=2) + "\n"
    )
    print(
        json.dumps(
            {
                "photos": len(results),
                "referenceMaxAbsErrors": [
                    r["referenceMaxAbsError"] for r in results
                ],
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path("vit-demo"))
    export(parser.parse_args().output_dir)