# examples/clip_export.py
"""Genuine CPU float32 official CLIP ViT-B/32 inference with saved full arrays.

Unzip clip_native.zip beside this file and follow README-clip.md. Model source
is unchanged, MIT-licensed, hash-authenticated native PyTorch. The tokenizer
adapter accepts printable ASCII only, rejects embedded special tokens, and uses
the exact pinned byte map, BPE merge slice and native merge algorithm. It does
not supply ftfy/regex Unicode normalization. Import performs no I/O or inference.
"""

import argparse
import ast
import gzip
import hashlib
import html
import importlib.util
import json
import platform
import re
import shutil
import types
from pathlib import Path
from urllib.request import urlopen

import numpy as np
import PIL
import torch
import torchvision
from PIL import Image
from torchvision import transforms

import clip_training

REVISION = "d05afc436d78f1c48dc0dbf8e5980a9d471f35f6"
CHECKPOINT_SHA256 = (
    "40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af"
)
CHECKPOINT_URL = (
    "https://openaipublic.azureedge.net/clip/models/"
    + CHECKPOINT_SHA256
    + "/ViT-B-32.pt"
)
MEAN = (0.48145466, 0.4578275, 0.40821073)
STD = (0.26862954, 0.26130258, 0.27577711)
CLASSES = ["cat", "dog", "cup of coffee", "bowl of soup", "car"]
TEMPLATES = ["a photo of a {}.", "a close-up photo of a {}."]
PHOTO_ROOT = "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.25.2/skimage/data/"
PHOTO_LICENSE = "https://github.com/scikit-image/scikit-image/blob/v0.25.2/skimage/data/_fetchers.py"
PHOTOS = [
    {
        "id": "cat",
        "title": "Chelsea the cat",
        "author": "Stefan van der Walt",
        "filename": "chelsea.png",
        "sha256": "596aa1e7cb875eb79f437e310381d26b338a81c2da23439704a73c4651e8c4bb",
    },
    {
        "id": "coffee",
        "title": "Cup of coffee",
        "author": "Rachel Michetti",
        "filename": "coffee.png",
        "sha256": "cc02f8ca188b167c775a7101b5d767d1e71792cf762c33d6fa15a4599b5a8de7",
    },
]


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


def authenticate_source(source_dir):
    """Validate revision and every native source against the public hash file."""
    source_dir = Path(source_dir)
    manifest = json.loads((source_dir / "source-manifest.json").read_text())
    expected = json.loads(
        Path(__file__).with_name("clip-source-hashes.json").read_text()
    )
    if manifest["revision"] != REVISION or manifest["files"] != expected:
        raise ValueError("Unexpected source revision or hash manifest.")
    for filename, record in expected.items():
        if sha256(source_dir / filename) != record["sha256"]:
            raise ValueError(f"Native source hash mismatch: {filename}")
    return manifest


def printable_ascii(text):
    """Reject Unicode and control characters, including tabs/newlines/DEL."""
    if not isinstance(text, str) or any(not 32 <= ord(c) <= 126 for c in text):
        raise ValueError("Tokenizer supports printable ASCII text only.")
    return text


class AsciiTokenizer:
    """Narrow native-compatible BPE adapter; no unrestricted ftfy parity claim."""

    def __init__(self, source_dir):
        source_dir = Path(source_dir)
        authenticate_source(source_dir)
        tree = ast.parse((source_dir / "clip/simple_tokenizer.py").read_text())
        definitions = [
            node
            for node in tree.body
            if isinstance(node, ast.FunctionDef)
            and node.name in ("bytes_to_unicode", "get_pairs")
        ]
        # Execute unchanged native helpers and bpe method, skipping optional imports.
        for node in definitions:
            node.decorator_list = []
        native_class = next(
            node for node in tree.body if isinstance(node, ast.ClassDef)
        )
        definitions.append(
            next(
                node
                for node in native_class.body
                if isinstance(node, ast.FunctionDef) and node.name == "bpe"
            )
        )
        namespace = {}
        exec(
            compile(
                ast.Module(body=definitions, type_ignores=[]),
                "native-bpe",
                "exec",
            ),
            namespace,
        )
        self.byte_encoder = namespace["bytes_to_unicode"]()
        self.bpe = types.MethodType(namespace["bpe"], self)
        lines = (
            gzip.decompress(
                (source_dir / "clip/bpe_simple_vocab_16e6.txt.gz").read_bytes()
            )
            .decode("utf8")
            .split("\n")
        )
        merges = [
            tuple(line.split()) for line in lines[1 : 49152 - 256 - 2 + 1]
        ]
        vocab = list(self.byte_encoder.values())
        vocab += [value + "</w>" for value in self.byte_encoder.values()]
        vocab += ["".join(pair) for pair in merges]
        vocab += ["<|startoftext|>", "<|endoftext|>"]
        self.encoder = dict(zip(vocab, range(len(vocab))))
        self.bpe_ranks = dict(zip(merges, range(len(merges))))
        self.cache = {
            "<|startoftext|>": "<|startoftext|>",
            "<|endoftext|>": "<|endoftext|>",
        }
        self.pattern = re.compile(
            r"'s|'t|'re|'ve|'m|'ll|'d|[a-z]+|[0-9]|[^\sa-z0-9]+", re.IGNORECASE
        )

    def encode(self, text):
        """Two HTML cleanup passes, spaces, lowercase, ASCII native token split."""
        text = printable_ascii(text)
        text = printable_ascii(html.unescape(html.unescape(text)))
        if html.unescape(text) != text:
            raise ValueError(
                "HTML entities must resolve within two cleanup passes."
            )
        text = re.sub(r" +", " ", text.strip()).lower()
        if "<|startoftext|>" in text or "<|endoftext|>" in text:
            raise ValueError(
                "Embedded special tokens are unsupported prompt text."
            )
        output = []
        for token in self.pattern.findall(text):
            encoded = "".join(
                self.byte_encoder[byte] for byte in token.encode("utf8")
            )
            output.extend(
                self.encoder[piece] for piece in self.bpe(encoded).split(" ")
            )
        return output

    def tokenize(self, texts):
        """Official 77-token context: SOT 49406, EOT 49407, zero padding."""
        return clip_training.pad_tokens(
            [self.encode(text) for text in texts],
            context=77,
            sot=49406,
            eot=49407,
        )


def load_native(source_dir, checkpoint):
    """Load full ViT-B/32 with strict native checkpoint and architecture checks."""
    authenticate_source(source_dir)
    if sha256(checkpoint) != CHECKPOINT_SHA256:
        raise ValueError("Official ViT-B/32 checkpoint hash mismatch.")
    spec = importlib.util.spec_from_file_location(
        "clip_native_model", Path(source_dir) / "clip/model.py"
    )
    native = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(native)
    archive = torch.jit.load(str(checkpoint), map_location="cpu").eval()
    state = archive.state_dict()
    del archive
    for key in ("input_resolution", "context_length", "vocab_size"):
        state.pop(key, None)
    model = native.CLIP(
        embed_dim=512,
        image_resolution=224,
        vision_layers=12,
        vision_width=768,
        vision_patch_size=32,
        context_length=77,
        vocab_size=49408,
        transformer_width=512,
        transformer_heads=8,
        transformer_layers=12,
    ).float()
    model.load_state_dict(state, strict=True)
    del state
    model.eval()
    for parameter in model.parameters():
        parameter.requires_grad_(False)
    if sum(p.numel() for p in model.parameters()) != 151277313:
        raise ValueError("Unexpected ViT-B/32 parameter count.")
    return model


def preprocess():
    """Exact official PIL bicubic resize, center crop, RGB and normalization."""
    return transforms.Compose(
        [
            transforms.Resize(
                224, interpolation=transforms.InterpolationMode.BICUBIC
            ),
            transforms.CenterCrop(224),
            transforms.Lambda(lambda image: image.convert("RGB")),
            transforms.ToTensor(),
            transforms.Normalize(MEAN, STD),
        ]
    )


def export(args):
    """Save full image/text features, tokens, prototypes and conditional outputs."""
    torch.set_num_threads(2)
    torch.manual_seed(12)
    output = args.output
    output.mkdir(parents=True, exist_ok=True)
    model = load_native(args.source_dir, args.checkpoint)
    tokenizer = AsciiTokenizer(args.source_dir)
    prompts = [
        template.format(label) for label in CLASSES for template in TEMPLATES
    ]
    tokens = tokenizer.tokenize(prompts)
    with torch.inference_mode():
        text = model.encode_text(tokens)
        normalized_text = clip_training.unit_features(text)
        prototypes = clip_training.class_prototypes(
            text.reshape(len(CLASSES), len(TEMPLATES), -1)
        )
    photo_records = []
    for photo in PHOTOS:
        original = output / f"source-{photo['id']}.png"
        if args.photo_dir:
            shutil.copyfile(
                args.photo_dir / f"source-{photo['id']}.png", original
            )
        else:
            with urlopen(PHOTO_ROOT + photo["filename"]) as response:
                original.write_bytes(response.read())
        if sha256(original) != photo["sha256"]:
            raise ValueError("Pinned photograph hash mismatch.")
        with Image.open(original) as image:
            original_size = [image.height, image.width]
            inputs = preprocess()(image).unsqueeze(0)
        crop = (
            inputs[0] * torch.tensor(STD)[:, None, None]
            + torch.tensor(MEAN)[:, None, None]
        )
        processed = output / f"input-{photo['id']}.png"
        transforms.ToPILImage()(crop.clamp(0, 1)).save(processed)
        with torch.inference_mode():
            image_features = model.encode_image(inputs)
            normalized_image = clip_training.unit_features(image_features)
            scores = clip_training.candidate_scores(
                image_features, prototypes, model.logit_scale
            )
            direct_logits = (
                model.logit_scale.exp() * normalized_image @ prototypes.T
            )
            candidate_parity = float(
                (scores["logits"] - direct_logits).abs().max()
            )
            if candidate_parity > 1e-5:
                raise ValueError(
                    "Candidate-logit native arithmetic parity failed."
                )
            native_image_logits, native_text_logits = model(inputs, tokens)
            reconstructed = (
                model.logit_scale.exp() * normalized_image @ normalized_text.T
            )
            parity = float((native_image_logits - reconstructed).abs().max())
            if parity > 1e-5 or not torch.equal(
                native_text_logits, native_image_logits.T
            ):
                raise ValueError("Native forward prompt-logit parity failed.")
        arrays = output / f"arrays-{photo['id']}.npz"
        np.savez_compressed(
            arrays,
            input=inputs.numpy(),
            tokens=tokens.numpy(),
            imageFeatures=image_features.numpy(),
            normalizedImage=normalized_image.numpy(),
            textFeatures=text.numpy(),
            normalizedText=normalized_text.numpy(),
            prototypes=prototypes.numpy(),
            promptLogits=native_image_logits.numpy(),
            candidateLogits=scores["logits"].numpy(),
            cosine=scores["cosine"].numpy(),
            softmax=scores["softmax"].numpy(),
            logitScale=model.logit_scale.numpy(),
        )
        photo_records.append(
            {
                "id": photo["id"],
                "title": photo["title"],
                "originalSize": original_size,
                "originalImage": original.name,
                "inputImage": processed.name,
                "inputImageSha256": sha256(processed),
                "arrays": arrays.name,
                "arraysSha256": sha256(arrays),
                "inputTensorSha256": hashlib.sha256(
                    inputs.numpy().tobytes()
                ).hexdigest(),
                "source": {
                    "url": PHOTO_ROOT + photo["filename"],
                    "sha256": photo["sha256"],
                    "author": photo["author"],
                    "license": "CC0",
                    "licenseUrl": PHOTO_LICENSE,
                },
                "cosine": scores["cosine"][0].tolist(),
                "logits": scores["logits"][0].tolist(),
                "softmax": scores["softmax"][0].tolist(),
                "nativePromptParityMaxError": parity,
                "nativeCandidateArithmeticMaxError": candidate_parity,
            }
        )
    manifest = {
        "scope": "genuine released official CLIP ViT-B/32; two photos, no benchmark",
        "provenance": {
            "model": "CLIP ViT-B/32",
            "sourceRevision": REVISION,
            "sourceHashes": authenticate_source(args.source_dir)["files"],
            "checkpointUrl": CHECKPOINT_URL,
            "checkpointSha256": CHECKPOINT_SHA256,
            "programSha256": sha256(clip_training.__file__),
            "exportSha256": sha256(__file__),
            "sourceArchiveSha256": sha256(
                Path(__file__).with_name("clip_native.zip")
            ),
            "python": platform.python_version(),
            "torch": torch.__version__,
            "torchvision": torchvision.__version__,
            "numpy": np.__version__,
            "Pillow": PIL.__version__,
            "device": "cpu",
            "dtype": "float32",
            "threads": 2,
            "parameterCount": sum(p.numel() for p in model.parameters()),
            "reproductionCommand": "python clip_export.py --source-dir clip_native --checkpoint ViT-B-32.pt --output clip-demo",
            "adapters": "Native model unchanged; native byte map and BPE method with printable ASCII token splitting/cleanup; no ftfy or regex normalization.",
        },
        "preprocessing": {
            "shortSide": 224,
            "centerCrop": [224, 224],
            "interpolation": "PIL bicubic",
            "rgbAfterCrop": True,
            "mean": MEAN,
            "std": STD,
        },
        "classes": CLASSES,
        "templates": TEMPLATES,
        "prompts": prompts,
        "tokens": tokens.tolist(),
        "contextLength": 77,
        "vocabSize": 49408,
        "logitScaleLog": float(model.logit_scale),
        "logitScaleExp": float(model.logit_scale.exp()),
        "scalePolicy": "raw released evaluation exp(logit_scale), no training cap",
        "conditionality": "Softmax is conditional on these candidate prototypes and templates; uncalibrated.",
        "photos": photo_records,
    }
    (output / "predictions.json").write_text(
        json.dumps(manifest, indent=2) + "\n"
    )
    print(
        json.dumps(
            {
                "photos": len(photo_records),
                "model": manifest["provenance"]["model"],
                "maxNativePromptParityError": max(
                    p["nativePromptParityMaxError"] for p in photo_records
                ),
            }
        )
    )


def main():
    """Require explicit local weights; retrieve only pinned licensed photos."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=Path(__file__).with_name("clip_native"),
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--photo-dir", type=Path)
    parser.add_argument("--output", type=Path, default=Path("clip-demo"))
    export(parser.parse_args())


if __name__ == "__main__":
    main()