# examples/byol.py
"""Offline BYOL update mechanics on tiny synthetic CPU images.

Authored PyTorch reference to paper 2006.07733v3, section 3.1 equations
1-3 and Appendix A. Update timing and target training-mode normalization
follow google-deepmind/deepmind-research revision
82a347438fd93bdd4bc01764258223ae583e5cbf, byol/byol_experiment.py
loss_fn and _update_fn; utils/helpers.py, networks.py, schedules.py.
This is not pretrained BYOL or an ImageNet reproduction. Linear encoders,
small BN MLPs, flip/noise views and SGD replace ResNet, full augmentations,
cross-replica BN and LARS. Copy initialization is a teaching choice; the
pinned release initializes its branches with different random keys.
"""

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

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


def validate_matrix(value, name):
    """Reject non-finite or empty floating point [B, D] CPU tensors."""
    if (not isinstance(value, torch.Tensor) or value.ndim != 2
            or min(value.shape) < 1 or not value.is_floating_point()
            or value.device.type != "cpu" or not torch.isfinite(value).all()):
        raise ValueError(f"{name} must be a finite floating CPU [B, D] tensor")


# region: views
def make_views(images, seed):
    """Make reproducible flip/noise views without changing global RNG state."""
    if (not isinstance(images, torch.Tensor) or images.ndim != 4
            or images.shape[0] < 2 or min(images.shape[1:]) < 1
            or not images.is_floating_point() or images.device.type != "cpu"
            or not torch.isfinite(images).all()):
        raise ValueError("images must be finite floating CPU [B,C,H,W], B >= 2")
    generator = torch.Generator().manual_seed(seed)
    noise1 = torch.randn(images.shape, generator=generator, dtype=images.dtype)
    noise2 = torch.randn(images.shape, generator=generator, dtype=images.dtype)
    return images + 0.05 * noise1, images.flip(-1) + 0.05 * noise2
# endregion: views


# region: branches
class TinyBranch(nn.Module):
    """Reduced encoder f and BN projector g with output [B, D]."""

    def __init__(self, input_dim=48, hidden_dim=12, projection_dim=6):
        super().__init__()
        if any(not isinstance(n, int) or n < 1
               for n in (input_dim, hidden_dim, projection_dim)):
            raise ValueError("branch dimensions must be positive integers")
        self.encoder = nn.Sequential(
            nn.Flatten(), nn.Linear(input_dim, hidden_dim), nn.ReLU())
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim),
            nn.ReLU(), nn.Linear(hidden_dim, projection_dim))

    def forward(self, images):
        """Project image features; training BN needs at least two samples."""
        return self.projector(self.encoder(images))


class TinyOnline(nn.Module):
    """Online f, g and predictor q; only f and g have target counterparts."""

    def __init__(self, input_dim=48, hidden_dim=12, projection_dim=6):
        super().__init__()
        self.branch = TinyBranch(input_dim, hidden_dim, projection_dim)
        self.predictor = nn.Sequential(
            nn.Linear(projection_dim, hidden_dim), nn.BatchNorm1d(hidden_dim),
            nn.ReLU(), nn.Linear(hidden_dim, projection_dim))

    def forward(self, images):
        """Return predictions [B, D] and projections [B, D]."""
        projection = self.branch(images)
        return self.predictor(projection), projection
# endregion: branches


# region: target
def make_target(online):
    """Copy f/g weights and buffers independently, then freeze parameters."""
    target = copy.deepcopy(online.branch)
    target.requires_grad_(False)
    target.train()  # Freeze weights, not running normalization state.
    return target


def target_forward(target, view1, view2):
    """Stop target graphs while updating BN state on each training view."""
    with torch.no_grad():
        z1 = target(view1)
        z2 = target(view2)
    return z1, z2
# endregion: target


# region: regression
def normalized_regression(prediction, target):
    """Return per-example 2 - 2 cosine, with no cross-example comparisons."""
    validate_matrix(prediction, "prediction")
    validate_matrix(target, "target")
    if prediction.shape != target.shape or prediction.dtype != target.dtype:
        raise ValueError("prediction and target shapes and dtypes must match")
    # Reject undefined directions rather than silently normalizing zero vectors.
    for value in (prediction, target):
        norms = torch.linalg.vector_norm(value, dim=-1)
        if not torch.isfinite(norms).all() or (norms <= 1e-6).any():
            raise ValueError("regression requires finite norms greater than 1e-6")
    p = F.normalize(prediction, dim=-1)
    z = F.normalize(target.detach(), dim=-1)
    return 2 - 2 * (p * z).sum(dim=-1)


def byol_loss(p1, p2, z1, z2):
    """Sum both view orderings before the batch mean, as paper section 3.1."""
    forward = normalized_regression(p1, z2)
    reverse = normalized_regression(p2, z1)
    if forward.shape != reverse.shape:
        raise ValueError("both view orderings must have the same batch size")
    return (forward + reverse).mean()
# endregion: regression


# region: online_update
def online_update(online, target, optimizer, view1, view2):
    """Backpropagate the symmetric loss and update only online parameters."""
    online_ids = {id(p) for p in online.parameters()}
    optimized_ids = {id(p) for group in optimizer.param_groups
                     for p in group["params"]}
    if optimized_ids != online_ids:
        raise ValueError("optimizer must contain exactly the online parameters")
    expected = online.branch.encoder[1].in_features
    for view in (view1, view2):
        if (view.ndim != 4 or view.shape[0] < 2
                or math.prod(view.shape[1:]) != expected):
            raise ValueError("views must be matching [B,C,H,W] with B >= 2")
        validate_matrix(view.flatten(1), "view")
    if view1.shape != view2.shape or view1.dtype != view2.dtype:
        raise ValueError("views must have matching shapes and dtypes")
    if view1.dtype != online.branch.encoder[1].weight.dtype:
        raise ValueError("view dtype must match online weights")
    if (not online.training or not target.training
            or any(p.requires_grad for p in target.parameters())):
        raise ValueError("train both branches with frozen target parameters")
    optimizer.zero_grad(set_to_none=True)
    p1, projection1 = online(view1)
    p2, _ = online(view2)
    z1, z2 = target_forward(target, view1, view2)
    loss = byol_loss(p1, p2, z1, z2)
    loss.backward()
    optimizer.step()  # Target EMA is a separate, subsequent operation.
    return loss.detach(), projection1.detach(), z1, z2
# endregion: online_update


# region: ema
def cosine_tau(step, total_steps, base_tau=0.996):
    """Increase tau from base_tau to one, including both endpoints."""
    if (not isinstance(step, int) or not isinstance(total_steps, int)
            or total_steps < 1 or not 0 <= step <= total_steps
            or not math.isfinite(base_tau) or not 0 <= base_tau <= 1):
        raise ValueError("invalid EMA schedule arguments")
    return 1 - (1 - base_tau) * (1 + math.cos(math.pi * step / total_steps)) / 2


def update_target(target, online_branch, tau):
    """Apply post-optimizer EMA to parameters, leaving BN buffers alone."""
    if not math.isfinite(tau) or not 0 <= tau <= 1:
        raise ValueError("tau must be finite and in [0, 1]")
    target_params = dict(target.named_parameters())
    online_params = dict(online_branch.named_parameters())
    if target_params.keys() != online_params.keys() or any(
            target_params[k].shape != online_params[k].shape
            for k in target_params):
        raise ValueError("target and online branch parameter structures differ")
    with torch.no_grad():
        for key, target_param in target_params.items():
            target_param.mul_(tau).add_(online_params[key], alpha=1 - tau)
    # BN running mean/variance already changed in target_forward, not via EMA.
# endregion: ema


def smoke(seed=7):
    """Execute one seeded CPU step and return measured finite diagnostics."""
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(seed)
        online = TinyOnline()
        target = make_target(online)
        images = torch.randn(4, 3, 4, 4)
        views = make_views(images, seed + 1)
        with torch.no_grad():
            representations = online.branch.encoder(views[0])
        old = {k: p.detach().clone() for k, p in target.named_parameters()}
        bn_before = target.projector[1].running_mean.clone()
        optimizer = torch.optim.SGD(online.parameters(), lr=0.05)
        loss, projection, z1, z2 = online_update(
            online, target, optimizer, *views)
        tau = cosine_tau(0, 10)
        grad_norms = {
            name: float(sum(p.grad.square().sum() for p in module.parameters()
                            if p.grad is not None).sqrt())
            for name, module in (("encoder", online.branch.encoder),
                                 ("projector", online.branch.projector),
                                 ("predictor", online.predictor))}
        bn_after_forward = target.projector[1].running_mean.clone()
        update_target(target, online.branch, tau)
        expected = {k: tau * old[k] + (1 - tau) * p.detach()
                    for k, p in online.branch.named_parameters()}
        ema_error = max(float((p - expected[k]).abs().max())
                        for k, p in target.named_parameters())
        with torch.no_grad():
            online.eval()
            p1, _ = online(views[0])
            p2, _ = online(views[1])
            scale_error = float((byol_loss(p1, p2, z1, z2)
                                - byol_loss(3 * p1, 2 * p2, 4 * z1, 5 * z2)).abs())
        result = {
            "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            "seed": seed, "device": "cpu", "boundary": "tiny synthetic mechanics",
            "reductions": {"encoder": "48->12", "projection": 6,
                           "optimizer": "SGD", "views": "flip and Gaussian noise",
                           "normalization": "local PyTorch BN, not cross-replica BN"},
            "shapes": {"images": list(images.shape),
                       "representations": list(representations.shape),
                       "projections": list(projection.shape)},
            "loss": float(loss), "gradient_norms": grad_norms,
            "target_has_grad": any(p.grad is not None for p in target.parameters()),
            "tau": tau, "ema_max_error": ema_error,
            "positive_scale_loss_error": scale_error,
            "representation_variance": float(representations.var(dim=0, unbiased=False).mean()),
            "projection_variance": float(projection.var(dim=0, unbiased=False).mean()),
            "target_bn_mean_delta": float((bn_after_forward - bn_before).abs().max()),
            "target_bn_batches": int(target.projector[1].num_batches_tracked),
            "ema_preserved_target_bn": torch.equal(
                target.projector[1].running_mean, bn_after_forward),
        }
    return result


def main():
    """Print actual smoke observations without datasets or checkpoints."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    print(json.dumps(smoke(args.seed), sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()