# examples/vicreg_training.py
"""VICReg released loss mechanics and separate genuine frozen-backbone export.

The offline smoke uses a random small CNN and SGD, not ImageNet pretraining.
The optional export strictly loads the official ResNet-50 backbone checkpoint.
"""

import argparse
import hashlib
import json
import platform
from pathlib import Path

from PIL import Image, ImageOps, __version__ as pillow_version
import torch
from torch import nn
from torch.nn import functional as F
import torchvision
from torchvision import models, transforms

REVISION = "4e12602fd495af83efd1631fbe82523e6db092e0"
CHECKPOINT_URL = "https://dl.fbaipublicfiles.com/vicreg/resnet50.pth"
CHECKPOINT_SHA256 = (
    "c843e7652491ff2f712734619fd74d85e5e92ac902615665ad3fd50dc6ada591"
)


# region: views
def two_views(images, size=32):
    """Sample two view batches with reduced crop/color/blur augmentation."""
    if len(images) < 2 or not isinstance(size, int) or size < 8:
        raise ValueError("Provide at least two images and size >= 8.")
    transform = transforms.Compose(
        [
            transforms.RandomResizedCrop(
                size,
                scale=(0.08, 1.0),
                interpolation=transforms.InterpolationMode.BICUBIC,
            ),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply(
                [transforms.ColorJitter(0.4, 0.4, 0.2, 0.1)], p=0.8
            ),
            transforms.RandomGrayscale(p=0.2),
            transforms.RandomApply(
                [transforms.GaussianBlur(3, (0.1, 2.0))], p=0.5
            ),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )
    first = torch.stack([transform(image.convert("RGB")) for image in images])
    second = torch.stack([transform(image.convert("RGB")) for image in images])
    return first, second


# endregion: views


# region: encoder
class VICReg(nn.Module):
    """Small shared encoder with the released three-linear-layer head pattern."""

    def __init__(self, representation_dim=32, projection_dim=16):
        super().__init__()
        if any(
            not isinstance(d, int) or d < 1
            for d in (representation_dim, projection_dim)
        ):
            raise ValueError("Feature dimensions must be positive integers.")
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, representation_dim, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.expander = nn.Sequential(
            nn.Linear(representation_dim, projection_dim),
            nn.BatchNorm1d(projection_dim),
            nn.ReLU(),
            nn.Linear(projection_dim, projection_dim),
            nn.BatchNorm1d(projection_dim),
            nn.ReLU(),
            nn.Linear(projection_dim, projection_dim, bias=False),
        )

    def forward(self, images):
        if (
            images.ndim != 4
            or images.shape[1] != 3
            or min(images.shape[2:]) < 8
        ):
            raise ValueError("Expected images [N, 3, H, W], H and W >= 8.")
        if self.training and images.shape[0] < 2:
            raise ValueError("Training BatchNorm needs at least two images.")
        h = self.encoder(images)
        return h, self.expander(h)


# endregion: encoder


def validate_embeddings(first, second):
    """Reject undefined sample statistics and mismatched branch tensors."""
    if (
        first.ndim != 2
        or first.shape != second.shape
        or first.shape[0] < 2
        or first.shape[1] < 1
    ):
        raise ValueError("Expected matching [N, D] tensors, N >= 2, D >= 1.")
    for tensor in (first, second):
        if (
            not torch.is_floating_point(tensor)
            or not torch.isfinite(tensor).all()
        ):
            raise ValueError(
                "Embeddings must be finite floating-point values."
            )
    if first.device != second.device or first.dtype != second.dtype:
        raise ValueError("Branches must share dtype and device.")


# region: invariance
def invariance_loss(first, second):
    """Mean squared difference over both N and D, as in the released code."""
    validate_embeddings(first, second)
    return F.mse_loss(first, second)


# endregion: invariance


# region: variance
def variance_loss(first, second, epsilon=1e-4):
    """Average branch hinge penalties with unbiased sample variance."""
    validate_embeddings(first, second)
    if not isinstance(epsilon, (int, float)) or not 0 < epsilon <= 1:
        raise ValueError("Epsilon must be in (0, 1].")
    first = first - first.mean(dim=0)
    second = second - second.mean(dim=0)
    std_first = torch.sqrt(first.var(dim=0, correction=1) + epsilon)
    std_second = torch.sqrt(second.var(dim=0, correction=1) + epsilon)
    return F.relu(1 - std_first).mean() / 2 + F.relu(1 - std_second).mean() / 2


# endregion: variance


# region: covariance
def off_diagonal(matrix):
    """Select every off-diagonal entry, including both symmetric halves."""
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError("Expected a square covariance matrix.")
    mask = ~torch.eye(matrix.shape[0], dtype=torch.bool, device=matrix.device)
    return matrix[mask]


def covariance_loss(first, second):
    """Sum both branch covariance penalties, dividing each by D, not D²."""
    validate_embeddings(first, second)
    n, d = first.shape
    first = first - first.mean(dim=0)
    second = second - second.mean(dim=0)
    cov_first = first.T @ first / (n - 1)
    cov_second = second.T @ second / (n - 1)
    return (
        off_diagonal(cov_first).square().sum()
        + off_diagonal(cov_second).square().sum()
    ) / d


# endregion: covariance


# region: total
def vicreg_loss(first, second):
    """Return released 25/25/1 total and individually inspectable terms."""
    invariance = invariance_loss(first, second)
    variance = variance_loss(first, second)
    covariance = covariance_loss(first, second)
    total = 25 * invariance + 25 * variance + covariance
    return total, {
        "invariance": invariance,
        "variance": variance,
        "covariance": covariance,
    }


# endregion: total


# region: update
def train_step(model, optimizer, first, second):
    """Backpropagate through both branches and update shared parameters."""
    if first.shape != second.shape or first.ndim != 4 or first.shape[0] < 2:
        raise ValueError("View batches must match with N >= 2.")
    model.train()
    optimizer.zero_grad(set_to_none=True)
    h_first, z_first = model(first)
    _, z_second = model(second)
    loss, terms = vicreg_loss(z_first, z_second)
    loss.backward()
    gradient_norm = torch.sqrt(
        sum(
            p.grad.square().sum()
            for p in model.parameters()
            if p.grad is not None
        )
    )
    optimizer.step()
    return {
        "loss": loss.item(),
        "terms": {name: value.item() for name, value in terms.items()},
        "gradientNorm": gradient_norm.item(),
        "hShape": list(h_first.shape),
        "zShape": list(z_first.shape),
    }


# endregion: update


# region: frozen
def load_frozen_backbone(checkpoint):
    """Strictly load all 318 official tensors, with no classifier or expander."""
    checkpoint = Path(checkpoint)
    if (
        hashlib.sha256(checkpoint.read_bytes()).hexdigest()
        != CHECKPOINT_SHA256
    ):
        raise ValueError("Official checkpoint SHA-256 mismatch.")
    backbone = models.resnet50(weights=None)
    backbone.fc = nn.Identity()
    backbone.load_state_dict(
        torch.load(checkpoint, map_location="cpu", weights_only=True),
        strict=True,
    )
    backbone.requires_grad_(False)
    return backbone.eval()


def frozen_features(backbone, images):
    """Apply the official evaluation transform; return unnormalized h."""
    if not images:
        raise ValueError("Provide at least one RGB image.")
    transform = transforms.Compose(
        [
            transforms.Resize(
                256,
                interpolation=transforms.InterpolationMode.BILINEAR,
                antialias=True,
            ),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )
    backbone.eval()
    backbone.requires_grad_(False)
    batch = torch.stack([transform(image.convert("RGB")) for image in images])
    with torch.no_grad():
        return backbone(batch), batch


# endregion: frozen


def smoke():
    """Run one synthetic update offline, without checkpoints or downloads."""
    torch.manual_seed(7)
    model = VICReg()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    result = train_step(
        model, optimizer, torch.rand(4, 3, 32, 32), torch.rand(4, 3, 32, 32)
    )
    result.update(
        scope="Random small CNN, synthetic inputs; loss mechanics only.",
        seed=7,
    )
    return result


def export_inspection(output, checkpoint, photo_root):
    """Export genuine pretrained features and separate random-weight mechanics."""
    torch.manual_seed(7)
    output, photo_root = Path(output), Path(photo_root)
    photos = json.loads((photo_root / "predictions.json").read_text())[
        "photos"
    ]
    images = []
    for photo in photos:
        source = photo_root / photo["originalImage"]
        if (
            hashlib.sha256(source.read_bytes()).hexdigest()
            != photo["source"]["sha256"]
        ):
            raise ValueError(f"Source image hash mismatch: {source}")
        with Image.open(source) as image:
            images.append(image.convert("RGB"))
    first, second = two_views(images)
    model = VICReg().train()
    with torch.no_grad():
        _, z_first = model(first)
        _, z_second = model(second)
        loss, terms = vicreg_loss(z_first, z_second)
    mechanics = {
        "scope": "Random small CNN; licensed photos; no training update.",
        "loss": loss.item(),
        "zShape": list(z_first.shape),
        "terms": {name: value.item() for name, value in terms.items()},
    }
    backbone = load_frozen_backbone(checkpoint)
    # Deterministic flip is an inspection perturbation, not validation augmentation.
    feature_images = images + [ImageOps.mirror(image) for image in images]
    features, batch = frozen_features(backbone, feature_images)
    unit = F.normalize(features, dim=1)
    output.mkdir(parents=True, exist_ok=True)
    saved_photos = []
    mean = torch.tensor([0.485, 0.456, 0.406])[:, None, None]
    std = torch.tensor([0.229, 0.224, 0.225])[:, None, None]
    for index, photo in enumerate(photos):
        inputs = []
        for label, i in [("original", index), ("flip", index + len(photos))]:
            filename = f'{photo["id"]}-{label}.png'
            transforms.ToPILImage()((batch[i] * std + mean).clamp(0, 1)).save(
                output / filename
            )
            inputs.append(
                {
                    "file": filename,
                    "sha256": hashlib.sha256(
                        (output / filename).read_bytes()
                    ).hexdigest(),
                }
            )
        saved_photos.append(
            {
                "id": photo["id"],
                "title": photo["title"],
                "source": photo["source"],
                "inputs": inputs,
            }
        )
    feature_file = output / "features.json"
    feature_file.write_text(json.dumps(features.tolist()) + "\n")
    result = {
        "scope": "Genuine frozen official VICReg ResNet-50 features; no classifier or benchmark evaluation.",
        "photos": saved_photos,
        "order": [f'{p["id"]}-original' for p in photos]
        + [f'{p["id"]}-flip' for p in photos],
        "featureShape": list(features.shape),
        "inputShape": list(batch.shape),
        "featureNorms": features.norm(dim=1).tolist(),
        "cosine": (unit @ unit.T).tolist(),
        "featuresSha256": hashlib.sha256(
            feature_file.read_bytes()
        ).hexdigest(),
        "mechanics": mechanics,
        "provenance": {
            "architecture": "Official ResNet-50 backbone; torchvision equivalent; fc Identity; 318 strictly loaded tensors",
            "officialRevision": REVISION,
            "source": "https://github.com/facebookresearch/vicreg",
            "checkpointUrl": CHECKPOINT_URL,
            "checkpointSha256": CHECKPOINT_SHA256,
            "preprocessing": "RGB PIL; short-side Resize(256), bilinear, antialias=True; CenterCrop(224); ToTensor [0,1]; Normalize mean [0.485,0.456,0.406], std [0.229,0.224,0.225]. Deterministic horizontal flip before preprocessing for second pair.",
            "evaluationSource": f"https://github.com/facebookresearch/vicreg/blob/{REVISION}/evaluate.py#L200-L210",
            "python": platform.python_version(),
            "torch": torch.__version__,
            "torchvision": torchvision.__version__,
            "pillow": pillow_version,
            "device": "cpu",
            "dtype": "float32",
            "threads": torch.get_num_threads(),
            "seed": 7,
            "programSha256": hashlib.sha256(
                Path(__file__).read_bytes()
            ).hexdigest(),
            "command": "python examples/vicreg_training.py --export public/examples/vicreg-demo --checkpoint .context/paper-series/weights/vicreg-resnet50.pth",
        },
    }
    (output / "inspection.json").write_text(
        json.dumps(result, indent=2) + "\n"
    )
    return result


def main():
    """Keep imports side-effect free; configure CPU execution only at entry."""
    parser = argparse.ArgumentParser(description=__doc__)
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("--smoke", action="store_true")
    action.add_argument("--export", type=Path)
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument(
        "--photo-root", type=Path, default=Path("public/examples/vit-demo")
    )
    args = parser.parse_args()
    if args.export and not args.checkpoint:
        parser.error("--export requires --checkpoint (offline local file).")
    torch.set_num_threads(2)
    result = (
        smoke()
        if args.smoke
        else export_inspection(args.export, args.checkpoint, args.photo_root)
    )
    print(json.dumps(result))


if __name__ == "__main__":
    main()