# examples/efficientnet_export.py
"""Export six real, CPU-float32 predictions and exact model input crops.

Run beside efficientnet.py with the pinned requirements. No training is run.
Cached checkpoint/source files stay outside the repository. Published assets
are content-addressed; a failed export preserves the previous manifest.
"""

import argparse
from datetime import datetime, timezone
import hashlib
import json
import math
import os
from pathlib import Path
import shutil
import tempfile
from urllib.parse import urlparse
from urllib.request import urlopen

from PIL import Image, __version__ as PILLOW_VERSION
import torch
import torchvision
from torchvision.models import efficientnet_b0, efficientnet_b4
from torchvision.transforms import functional as F

from efficientnet import load_pretrained, prepare_image

HERE = Path(__file__).resolve().parent
REVISION = "78839c2b06c83c6cfb5c4da692ffb331bbd4c4cc"
PRESET_REVISION = "a2d5e429b9c57e0e63e6230f05046e4c0e4318f7"
SOURCE_BASE = "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.25.2/skimage/data/"
LICENSE_URL = "https://github.com/scikit-image/scikit-image/blob/v0.25.2/skimage/data/_fetchers.py"
PHOTOS = [
    ("cat", "Chelsea the cat", "chelsea.png", "CC0", "Stefan van der Walt"),
    ("coffee", "Coffee cup", "coffee.png", "CC0", "Rachel Michetti"),
    (
        "astronaut",
        "Astronaut Eileen Collins",
        "astronaut.png",
        "Public domain",
        "NASA",
    ),
]


def digest(path):
    """Hash exact file bytes."""
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def read_source(path):
    """Read a local RGB image or reject unreadable data."""
    try:
        with Image.open(path) as source:
            return source.convert("RGB")
    except (OSError, ValueError) as error:
        raise ValueError(f"Unreadable source image: {path}") from error


def publish_artifacts(staging, output_dir, manifest):
    """Publish immutable image bytes, then atomically replace the manifest."""
    for source in (staging / "images").rglob("*.png"):
        destination = output_dir / source.relative_to(staging)
        destination.parent.mkdir(parents=True, exist_ok=True)
        if destination.exists():
            if digest(destination) != digest(source):
                raise ValueError(
                    f"Content-addressed image collision: {destination}"
                )
        else:
            shutil.copyfile(source, destination)
    destination = output_dir / "data/efficientnet/predictions.json"
    destination.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary = tempfile.mkstemp(
        dir=destination.parent, suffix=".json"
    )
    try:
        with os.fdopen(descriptor, "w") as stream:
            json.dump(manifest, stream, indent=2, allow_nan=False)
            stream.write("\n")
        os.replace(temporary, destination)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def trace_predictions(model, batch):
    """Record actual intermediate shapes; always remove diagnostic hooks."""
    traces, handles = [], []

    def hook(name):
        def capture(_module, inputs, output):
            traces.append(
                {
                    "id": name,
                    "inputShape": list(inputs[0].shape),
                    "outputShape": list(output.shape),
                }
            )

        return capture

    try:
        for i, layer in enumerate(model.features):
            name = "stem" if i == 0 else "head" if i == 8 else f"stage-{i}"
            handles.append(layer.register_forward_hook(hook(name)))
        handles.append(model.avgpool.register_forward_hook(hook("pool")))
        handles.append(
            model.classifier.register_forward_hook(hook("classifier"))
        )
        with torch.inference_mode():
            logits = model(batch)
        return logits, traces
    finally:
        for handle in handles:
            handle.remove()


def export(output_dir, num_threads=2):
    """Strictly load both classifiers, verify each run, and publish six cases."""
    if num_threads < 1:
        raise ValueError("Thread count must be positive")
    for actual, expected in [
        (torch.__version__, "2.12.0"),
        (torchvision.__version__, "0.27.0"),
        (PILLOW_VERSION, "12.2.0"),
    ]:
        if actual.split("+")[0] != expected:
            raise ValueError(f"Expected version {expected}; got {actual}")
    torch.set_num_threads(num_threads)
    cache = Path(torch.hub.get_dir()) / "efficientnet-samples"
    cache.mkdir(parents=True, exist_ok=True)
    manifest = {
        "schemaVersion": 1,
        "provenance": {
            "sourceRevision": REVISION,
            "presetRevision": PRESET_REVISION,
            "referenceSha256": digest(HERE / "efficientnet.py"),
            "exportSha256": digest(Path(__file__)),
            "versions": {
                "torch": torch.__version__,
                "torchvision": torchvision.__version__,
                "Pillow": PILLOW_VERSION,
            },
            "device": "cpu",
            "dtype": "float32",
            "numThreads": num_threads,
            "exportedAt": datetime.now(timezone.utc).isoformat(),
        },
        "categories": [],
        "models": [],
        "photos": [],
        "runs": [],
    }
    with tempfile.TemporaryDirectory(
        prefix="efficientnet-export-"
    ) as directory:
        staging = Path(directory)
        assets = staging / "images/papers/efficientnet"
        assets.mkdir(parents=True)
        images = {}
        for photo_id, label, filename, license_name, attribution in PHOTOS:
            source = cache / filename
            if not source.exists():
                with urlopen(SOURCE_BASE + filename, timeout=60) as response:
                    source.write_bytes(response.read())
            image = read_source(source)
            source_hash = digest(source)
            saved = assets / f"source-{photo_id}-{source_hash[:12]}.png"
            shutil.copyfile(source, saved)
            images[photo_id] = image
            manifest["photos"].append(
                {
                    "id": photo_id,
                    "label": label,
                    "originalImage": "/" + str(saved.relative_to(staging)),
                    "originalShape": [image.height, image.width],
                    "source": {
                        "url": SOURCE_BASE + filename,
                        "sha256": source_hash,
                        "license": license_name,
                        "licenseUrl": LICENSE_URL,
                        "attribution": attribution,
                    },
                }
            )
        for model_id, builder in [
            ("b0", efficientnet_b0),
            ("b4", efficientnet_b4),
        ]:
            model, weights = load_pretrained(model_id)
            reference = builder(weights=None).eval()
            reference.load_state_dict(model.state_dict(), strict=True)
            checkpoint = (
                Path(torch.hub.get_dir())
                / "checkpoints"
                / Path(urlparse(weights.url).path).name
            )
            transform = weights.transforms()
            if (
                manifest["categories"]
                and manifest["categories"] != weights.meta["categories"]
            ):
                raise ValueError("Checkpoint categories differ")
            manifest["categories"] = weights.meta["categories"]
            manifest["models"].append(
                {
                    "id": model_id,
                    "label": f"EfficientNet-{model_id.upper()}",
                    "parameterCount": sum(
                        p.numel() for p in model.parameters()
                    ),
                    "strictLoad": True,
                    "weights": {
                        "enum": type(weights).__name__ + "." + weights.name,
                        "url": weights.url,
                        "sha256": digest(checkpoint),
                        "bytes": checkpoint.stat().st_size,
                    },
                    "transform": {
                        "resizeSize": transform.resize_size[0],
                        "cropSize": transform.crop_size[0],
                        "interpolation": "bicubic",
                        "antialias": transform.antialias,
                        "mean": transform.mean,
                        "std": transform.std,
                    },
                    "metadata": {
                        "sourceLabel": "torchvision 0.27.0 checkpoint metadata",
                        "acc1": weights.meta["_metrics"]["ImageNet-1K"][
                            "acc@1"
                        ],
                        "acc5": weights.meta["_metrics"]["ImageNet-1K"][
                            "acc@5"
                        ],
                        "operations": weights.meta["_ops"],
                        "operationsLabel": "torchvision GFLOPS metadata",
                    },
                }
            )
            for photo_id, image in images.items():
                crop, batch, record = prepare_image(image, weights)
                torch.testing.assert_close(
                    batch[0], transform(image), rtol=0, atol=0
                )
                temporary = assets / "input.png"
                crop.save(temporary)
                crop_hash = digest(temporary)
                saved = assets / f"{photo_id}-{model_id}-{crop_hash[:12]}.png"
                temporary.rename(saved)
                restored = F.convert_image_dtype(
                    F.pil_to_tensor(read_source(saved)), torch.float32
                )
                restored = F.normalize(
                    restored, transform.mean, transform.std
                ).unsqueeze(0)
                if not torch.equal(restored, batch):
                    raise ValueError("Saved crop differs from input tensor")
                logits, traces = trace_predictions(model, batch)
                with torch.inference_mode():
                    expected = reference(batch)
                torch.testing.assert_close(
                    logits, expected, rtol=1e-5, atol=1e-6
                )
                probabilities = logits.softmax(dim=1)[0].tolist()
                raw = logits[0].tolist()
                if len(raw) != 1000 or not all(
                    math.isfinite(v) for v in raw + probabilities
                ):
                    raise ValueError("Invalid predictions")
                if abs(sum(probabilities) - 1) > 1e-6:
                    raise ValueError("Invalid probability sum")
                top5 = sorted(
                    range(1000), key=lambda i: (-probabilities[i], i)
                )[:5]
                manifest["runs"].append(
                    {
                        "id": f"{photo_id}-{model_id}",
                        "photoId": photo_id,
                        "modelId": model_id,
                        "inputImage": "/" + str(saved.relative_to(staging)),
                        "inputSha256": crop_hash,
                        "inputShape": list(batch.shape),
                        "resizedShape": record["resizedShape"],
                        "cropBox": record["cropBox"],
                        "logits": raw,
                        "probabilities": probabilities,
                        "top5": top5,
                        "stages": traces,
                        "verification": {
                            "inputExact": True,
                            "referenceClose": True,
                            "maxAbsoluteLogitError": (logits - expected)
                            .abs()
                            .max()
                            .item(),
                            "rtol": 1e-5,
                            "atol": 1e-6,
                        },
                    }
                )
                print(
                    f"Verified {photo_id}-{model_id}: {weights.meta['categories'][top5[0]]}",
                    flush=True,
                )
        expected_pairs = {
            (photo[0], model) for photo in PHOTOS for model in ("b0", "b4")
        }
        if {
            (run["photoId"], run["modelId"]) for run in manifest["runs"]
        } != expected_pairs:
            raise ValueError("Incomplete export")
        publish_artifacts(staging, output_dir, manifest)
    return manifest


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path("public"))
    parser.add_argument("--num-threads", type=int, default=2)
    args = parser.parse_args()
    export(args.output_dir, args.num_threads)