# examples/siglip.py
"""Offline CPU reference for original SigLIP all-pairs loss, not SigLIP2.

Authored PyTorch mechanics following paper Algorithm 1 and sections 3.2-3.3:
https://arxiv.org/html/2303.15343v2
Official big_vision loss/model convention pinned to:
642c086e68b5e49ec8c48cead8b1c6fb042fc85a
No pretrained towers, distributed collectives, or benchmark measurements.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F

UPSTREAM_REVISION = "642c086e68b5e49ec8c48cead8b1c6fb042fc85a"
NORM_EPS = 1e-8


def _positive_int(value, name):
    """Reject empty dimensions, bools, and nonintegral chunk/batch sizes."""
    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
        raise ValueError(f"{name} must be a positive integer")


def _features(image, text, paired=False):
    """Validate finite CPU float32/64 matrices with common channel width."""
    for value in (image, text):
        if not isinstance(value, torch.Tensor) or value.ndim != 2:
            raise ValueError("features must be [N, D] tensors")
        if min(value.shape) < 1 or value.dtype not in (torch.float32, torch.float64):
            raise ValueError("features must be nonempty float32/64 matrices")
        if value.device.type != "cpu" or not torch.isfinite(value).all():
            raise ValueError("features must be finite CPU tensors")
    if image.dtype != text.dtype or image.shape[1] != text.shape[1]:
        raise ValueError("features need a common dtype and channel width")
    if paired and image.shape[0] != text.shape[0]:
        raise ValueError("paired image and text batches must have equal length")


def _parameters(log_scale, bias, reference):
    """Validate two scalar tensors without detaching their autograd paths."""
    for value in (log_scale, bias):
        if not isinstance(value, torch.Tensor) or value.ndim != 0:
            raise ValueError("log scale and bias must be scalar tensors")
        if value.dtype != reference.dtype or value.device != reference.device:
            raise ValueError("parameters must match feature dtype and device")
        if not torch.isfinite(value):
            raise ValueError("parameters must be finite")
    scale = log_scale.exp()
    if not torch.isfinite(scale) or scale <= 0:
        raise ValueError("exponentiated log scale must be finite and positive")


# region: normalize
def normalize_features(image, text):
    """Normalize paired [N, D] features; zero rows stay zero with epsilon."""
    _features(image, text, paired=True)
    image_norm = torch.linalg.vector_norm(image, dim=-1, keepdim=True)
    text_norm = torch.linalg.vector_norm(text, dim=-1, keepdim=True)
    if not torch.isfinite(image_norm).all() or not torch.isfinite(text_norm).all():
        raise ValueError("feature norms must be finite")
    # Match the pinned two_towers.py norm + 1e-8 convention.
    return image / (image_norm + NORM_EPS), text / (text_norm + NORM_EPS)
# endregion: normalize


# region: logits
def pair_logits(image, text, log_scale, bias):
    """Map normalized [R, D], [C, D] features to logits [R, C]."""
    _features(image, text)
    _parameters(log_scale, bias, image)
    logits = log_scale.exp() * (image @ text.T) + bias
    if not torch.isfinite(logits).all():
        raise ValueError("computed logits must be finite")
    return logits
# endregion: logits


# region: signed_labels
def signed_labels(image_indices, text_indices, dtype=torch.float64):
    """Use global indices, including chunks with no positive pairs."""
    for indices in (image_indices, text_indices):
        if (not isinstance(indices, torch.Tensor) or indices.ndim != 1
                or indices.numel() == 0 or indices.dtype != torch.int64
                or indices.device.type != "cpu"):
            raise ValueError("indices must be nonempty CPU int64 vectors")
        if (indices < 0).any() or indices.unique().numel() != indices.numel():
            raise ValueError("indices must be nonnegative and unique")
    if dtype not in (torch.float32, torch.float64):
        raise ValueError("label dtype must be float32/64")
    matches = image_indices[:, None] == text_indices[None, :]
    return 2 * matches.to(dtype) - 1
# endregion: signed_labels


# region: stable_loss
def stable_loss(logits, labels, batch_size):
    """Sum stable signed pair losses and divide by global N, never N squared."""
    _positive_int(batch_size, "batch_size")
    if (not isinstance(logits, torch.Tensor) or logits.ndim != 2
            or min(logits.shape) < 1
            or logits.dtype not in (torch.float32, torch.float64)
            or logits.device.type != "cpu" or not torch.isfinite(logits).all()):
        raise ValueError("logits must be a nonempty finite CPU float matrix")
    if (not isinstance(labels, torch.Tensor) or labels.shape != logits.shape
            or labels.dtype != logits.dtype or labels.device != logits.device
            or not ((labels == 1) | (labels == -1)).all()):
        raise ValueError("labels must match logits and contain only -1 or +1")
    return F.softplus(-labels * logits).sum() / batch_size
# endregion: stable_loss


def dense_loss(image, text, log_scale, bias):
    """Compute the original objective for every paired minibatch combination."""
    zimage, ztext = normalize_features(image, text)
    n = image.shape[0]
    indices = torch.arange(n)
    labels = signed_labels(indices, indices, image.dtype)
    return stable_loss(pair_logits(zimage, ztext, log_scale, bias), labels, n)


# region: chunked_loss
def chunked_loss(image, text, log_scale, bias, chunk_size):
    """Accumulate global-index chunks with the same total sum/N reduction."""
    _positive_int(chunk_size, "chunk_size")
    zimage, ztext = normalize_features(image, text)
    n = image.shape[0]
    total = image.new_zeros(())
    for i in range(0, n, chunk_size):
        rows = torch.arange(i, min(i + chunk_size, n))
        for j in range(0, n, chunk_size):
            cols = torch.arange(j, min(j + chunk_size, n))
            logits = pair_logits(zimage[rows], ztext[cols], log_scale, bias)
            labels = signed_labels(rows, cols, image.dtype)
            total = total + stable_loss(logits, labels, n)
    return total
# endregion: chunked_loss


class SigLIPLoss(nn.Module):
    """Trainable log scale and bias with explicit recipe initialization.

    The configured scale is 10, so the stored log scale is log(10), not 10.
    Bias -10 is a recipe choice, not an optimal prior for every small N.
    """

    def __init__(self, scale_init=10.0, bias_init=-10.0, dtype=torch.float64):
        super().__init__()
        if (not math.isfinite(scale_init) or scale_init <= 0
                or not math.isfinite(bias_init)
                or dtype not in (torch.float32, torch.float64)):
            raise ValueError("initial scale must be positive; parameters finite floats")
        self.log_scale = nn.Parameter(torch.tensor(math.log(scale_init), dtype=dtype))
        self.bias = nn.Parameter(torch.tensor(bias_init, dtype=dtype))
        _parameters(self.log_scale, self.bias, self.bias)

    def forward(self, image, text, chunk_size=None):
        """Return dense or serial chunked all-pairs loss with live parameters."""
        if chunk_size is None:
            return dense_loss(image, text, self.log_scale, self.bias)
        return chunked_loss(image, text, self.log_scale, self.bias, chunk_size)


def smoke(seed=7):
    """Execute finite seeded CPU values, derivatives, and chunk equivalence."""
    generator = torch.Generator(device="cpu").manual_seed(seed)
    image = torch.randn(5, 4, generator=generator, dtype=torch.float64,
                        requires_grad=True)
    text = torch.randn(5, 4, generator=generator, dtype=torch.float64,
                       requires_grad=True)
    objective = SigLIPLoss()
    parameters = (image, text, objective.log_scale, objective.bias)
    loss = objective(image, text)
    gradients = torch.autograd.grad(loss, parameters)
    chunks = []
    for size in (1, 2, 3, 8):
        chunk_loss = objective(image, text, chunk_size=size)
        chunk_gradients = torch.autograd.grad(chunk_loss, parameters)
        chunks.append({
            "chunk_size": size,
            "loss": chunk_loss.item(),
            "value_abs_error": abs(chunk_loss.item() - loss.item()),
            "gradient_max_abs_error": max(
                (left - right).abs().max().item()
                for left, right in zip(gradients, chunk_gradients)),
        })
    extreme = torch.tensor([[1000.0, -1000.0], [-1000.0, 1000.0]],
                           dtype=torch.float64, requires_grad=True)
    extreme_loss = stable_loss(extreme, -torch.ones_like(extreme), 2)
    extreme_gradient = torch.autograd.grad(extreme_loss, extreme)[0]
    record = {
        "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "upstream_revision": UPSTREAM_REVISION,
        "seed": seed,
        "device": "cpu",
        "dtype": "float64",
        "boundary": "synthetic features; no pretrained towers or distributed performance",
        "reductions": {"batch": 5, "width": 4, "pair_count": 25,
                       "normalization": "sum all pair losses / N", "optimizer_steps": 0},
        "shapes": {"image": list(image.shape), "text": list(text.shape),
                   "dense_logits": [5, 5], "off_diagonal_chunk": [2, 2]},
        "initialization": {"scale": objective.log_scale.exp().item(),
                           "log_scale": objective.log_scale.item(),
                           "bias": objective.bias.item()},
        "dense_loss": loss.item(),
        "gradient_norms": dict(zip(("image", "text", "log_scale", "bias"),
                                   (value.norm().item() for value in gradients))),
        "chunks": chunks,
        "off_diagonal_positive_count": int((signed_labels(
            torch.tensor([0, 1]), torch.tensor([2, 3])) == 1).sum()),
        "extreme_loss": extreme_loss.item(),
        "extreme_gradient": extreme_gradient.tolist(),
        "finite": bool(torch.isfinite(loss) and all(
            torch.isfinite(value).all() for value in gradients)
            and torch.isfinite(extreme_loss) and torch.isfinite(extreme_gradient).all()
            and all(math.isfinite(item["loss"]) and math.isfinite(
                item["gradient_max_abs_error"]) for item in chunks)),
    }
    if not record["finite"]:
        raise RuntimeError("smoke produced a nonfinite observation")
    return record


def main():
    """Print actual JSON only when explicitly invoked with --smoke."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    arguments = parser.parse_args()
    if not arguments.smoke:
        parser.error("pass --smoke to execute the offline reference")
    print(json.dumps(smoke(arguments.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()