"""
SimCLR, mostly generated by Claude Opus 4

Paper "A Simple Framework for Contrastive Learning of Visual Representations",
https://arxiv.org/abs/2002.05709
"""

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


# pylint: disable=abstract-method,arguments-differ
class FullGatherLayer(torch.autograd.Function):
    @staticmethod
    def forward(_ctx, x):  # type: ignore
        if training_utils.is_dist_available_and_initialized() is False:
            return (x,)

        output = [torch.zeros_like(x) for _ in range(dist.get_world_size())]
        dist.all_gather(output, x)
        return tuple(output)

    @staticmethod
    def backward(_ctx, *grads):  # type: ignore
        all_gradients = torch.stack(grads)
        if training_utils.is_dist_available_and_initialized() is False:
            return all_gradients

        dist.all_reduce(all_gradients)
        return all_gradients[dist.get_rank()]


class SimCLR(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"]
        temperature: float = self.config["temperature"]

        self.temperature = temperature
        self.projection_head = nn.Sequential(
            nn.Linear(self.backbone.embedding_size, projection_hidden_dim),
            nn.ReLU(),
            nn.Linear(projection_hidden_dim, projection_dim, bias=False),
        )

    def _contrastive_loss(self, features: torch.Tensor) -> torch.Tensor:
        """
        Compute NT-Xent loss (Normalized Temperature-scaled Cross Entropy Loss)
        """

        device = features.device
        batch_size = features.size(0) // 2

        features = F.normalize(features, dim=1)
        similarity_matrix = torch.matmul(features, features.t())
        mask = torch.eye(batch_size * 2, dtype=torch.bool, device=device)

        similarity_matrix = similarity_matrix[~mask].reshape(batch_size * 2, batch_size * 2 - 1)
        similarity_matrix = similarity_matrix / self.temperature

        # Create labels: for sample i, positive is at i+batch_size (and vice versa)
        labels = torch.concat([torch.arange(batch_size) + batch_size - 1, torch.arange(batch_size)], dim=0)
        labels = labels.to(device)

        # Compute loss
        loss = F.cross_entropy(similarity_matrix, labels, reduction="mean")

        return loss

    # pylint: disable=arguments-differ
    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        h1 = self.backbone.embedding(x1)
        h2 = self.backbone.embedding(x2)

        z1 = self.projection_head(h1)
        z2 = self.projection_head(h2)

        features = torch.concat([z1, z2], dim=0)
        features = torch.concat(FullGatherLayer.apply(features), dim=0)

        loss = self._contrastive_loss(features)

        return loss
