"""
LeJEPA, adapted from
https://github.com/facebookresearch/eb_jepa/blob/main/eb_jepa/losses.py
and (PR #2665 by rwightman)
https://github.com/huggingface/pytorch-image-models/blob/6f1005af19cf7b5ab9e0b489cad43b0c7643d53b/timm/task/lejepa.py

Paper "LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics",
https://arxiv.org/abs/2511.08544

Generated by gpt-5.3-codex xhigh.
"""

# Reference license: Apache-2.0 (both)

from typing import Any
from typing import Optional

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

from birder.common import training_utils
from birder.net.base import BaseNet
from birder.net.ssl.base import SSLBaseNet


class SIGReg(nn.Module):
    """
    Sketched Isotropic Gaussian Regularization (SIGReg)

    Computes the Epps-Pulley-style discrepancy between the empirical characteristic
    function of random 1D projections and that of a standard Gaussian.

    Implements the SIGReg loss following the LeJEPA formulation, with
    codebase-specific implementation details.

    Paper references:
    - Algorithm 1 (SIGReg)
    - Algorithm 2 (LeJEPA objective composition)
    - Section 5.1 (loss definition)
    """

    def __init__(self, num_knots: int, num_slices: int, t_max: float) -> None:
        super().__init__()
        self.num_slices = num_slices

        t = torch.linspace(0.0, t_max, num_knots, dtype=torch.float32)
        dt = t_max / (num_knots - 1)

        weights = torch.full((num_knots,), 2.0 * dt, dtype=torch.float32)
        weights[[0, -1]] = dt

        phi_gaussian = torch.exp(-t.square() / 2.0)

        self.t = nn.Buffer(t)
        self.phi_gaussian = nn.Buffer(phi_gaussian)
        self.weights = nn.Buffer(weights * phi_gaussian)

    def forward(self, proj: torch.Tensor) -> torch.Tensor:
        # Compute in float32
        proj = proj.float()

        directions = torch.randn(proj.size(-1), self.num_slices, device=proj.device, dtype=proj.dtype)
        directions = F.normalize(directions, dim=0)

        # [V, B, num_slices]
        x_proj = proj @ directions

        # [V, B, num_slices, num_knots]
        x_t = x_proj.unsqueeze(-1) * self.t
        cos_mean = x_t.cos().mean(dim=-3)
        sin_mean = x_t.sin().mean(dim=-3)

        if training_utils.is_dist_available_and_initialized() is True:
            world_size = dist.get_world_size()
            cos_mean = dist.nn.all_reduce(cos_mean) / world_size
            sin_mean = dist.nn.all_reduce(sin_mean) / world_size
        else:
            world_size = 1

        err = (cos_mean - self.phi_gaussian).square() + sin_mean.square()
        statistic = (err @ self.weights) * (proj.size(-2) * world_size)

        return statistic.mean()


class LeJEPA(SSLBaseNet):
    def __init__(
        self,
        backbone: BaseNet,
        *,
        config: Optional[dict[str, Any]] = None,
        size: Optional[tuple[int, int]] = None,
    ) -> None:
        super().__init__(backbone, config=config, size=size)
        assert self.config is not None, "must set config"

        projection_dim: int = self.config["projection_dim"]
        projection_hidden_dim: int = self.config["projection_hidden_dim"]
        projection_layers: int = self.config["projection_layers"]
        num_knots: int = self.config.get("num_knots", 17)
        num_slices: int = self.config.get("num_slices", 256)
        t_max: float = self.config.get("t_max", 3.0)
        self.loss_lambda: float = self.config.get("loss_lambda", 0.02)

        layers: list[nn.Module] = []
        in_dim = self.backbone.embedding_size
        for _ in range(projection_layers):
            layers.extend(
                [
                    nn.Linear(in_dim, projection_hidden_dim),
                    nn.BatchNorm1d(projection_hidden_dim),
                    nn.ReLU(),
                ]
            )
            in_dim = projection_hidden_dim

        layers.append(nn.Linear(in_dim, projection_dim))
        self.projector = nn.Sequential(*layers)
        self.sigreg = SIGReg(num_knots=num_knots, num_slices=num_slices, t_max=t_max)

    def forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        projections: list[torch.Tensor] = []
        for view in x:
            embedding = self.backbone.embedding(view)
            projections.append(self.projector(embedding))

        proj = torch.stack(projections, dim=0)  # [V, B, D]

        sigreg_loss = self.sigreg(proj)
        inv_loss = (proj.float().mean(dim=0, keepdim=True) - proj.float()).square().mean()
        total_loss = self.loss_lambda * sigreg_loss + (1.0 - self.loss_lambda) * inv_loss

        return (total_loss, sigreg_loss, inv_loss)
