from __future__ import annotations

import torch
import torch.nn.functional as F


def off_diagonal_mask(size: int, device: torch.device | str) -> torch.Tensor:
    return ~torch.eye(size, dtype=torch.bool, device=device)


def standardized_relation(
    features: torch.Tensor, mask: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """Cosine relation field and its standardized off-diagonal values."""
    features = F.normalize(features.float(), dim=-1)
    relation = features @ features.T
    if mask is None:
        mask = off_diagonal_mask(len(features), features.device)
    values = relation[mask]
    standardized = (values - values.mean()) / values.std().clamp_min(1e-6)
    return relation, standardized


def relation_field_energy(
    visual_relation: torch.Tensor,
    visual_standardized: torch.Tensor,
    language_particles: torch.Tensor,
    temperatures: tuple[float, ...] = (0.03, 0.07, 0.15),
) -> tuple[torch.Tensor, torch.Tensor]:
    """Second-order and multiscale conditional relation energies."""
    language_relation, language_standardized = standardized_relation(
        language_particles
    )
    mse = F.mse_loss(language_standardized, visual_standardized)
    diagonal = torch.eye(
        len(language_particles),
        dtype=torch.bool,
        device=language_particles.device,
    )
    conditional_kl = language_particles.new_zeros(())
    for temperature in temperatures:
        visual_logits = (visual_relation / temperature).masked_fill(
            diagonal, -1e4
        )
        language_logits = (language_relation / temperature).masked_fill(
            diagonal, -1e4
        )
        visual_probability = F.softmax(visual_logits, dim=-1)
        conditional_kl = conditional_kl + (
            visual_probability
            * (
                F.log_softmax(visual_logits, dim=-1)
                - F.log_softmax(language_logits, dim=-1)
            )
        ).sum(-1).mean()
    return mse, conditional_kl


def projection_quantile_target(
    text_features: torch.Tensor,
    particles: int,
    projections: int,
    generator: torch.Generator,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fixed sliced-distribution target from an unpaired text population."""
    directions = F.normalize(
        torch.randn(
            text_features.shape[-1],
            projections,
            generator=generator,
            device=text_features.device,
        ),
        dim=0,
    )
    projected = (text_features @ directions).sort(dim=0).values
    quantile_indices = (
        torch.linspace(
            0,
            len(projected) - 1,
            particles,
            device=text_features.device,
        )
        .round()
        .long()
    )
    return directions, projected[quantile_indices]


def sliced_distribution_energy(
    particles: torch.Tensor,
    directions: torch.Tensor,
    target_quantiles: torch.Tensor,
) -> torch.Tensor:
    projected = (F.normalize(particles, dim=-1) @ directions).sort(
        dim=0
    ).values
    return F.mse_loss(projected, target_quantiles)


def prototype_manifold_energy(
    particles: torch.Tensor, prototypes: torch.Tensor
) -> torch.Tensor:
    particles = F.normalize(particles, dim=-1)
    prototypes = F.normalize(prototypes, dim=-1)
    return (1 - (particles @ prototypes.T).max(dim=-1).values).mean()


def log_sinkhorn(
    logits: torch.Tensor, temperature: float, iterations: int = 12
) -> torch.Tensor:
    """Doubly stochastic coupling with differentiable log-domain updates."""
    log_coupling = logits / temperature
    for _ in range(iterations):
        log_coupling = log_coupling - torch.logsumexp(
            log_coupling, dim=1, keepdim=True
        )
        log_coupling = log_coupling - torch.logsumexp(
            log_coupling, dim=0, keepdim=True
        )
    return log_coupling.exp()


def retrieval_metrics(
    particles: torch.Tensor, paired_text: torch.Tensor
) -> dict[str, float]:
    particles = F.normalize(particles.float(), dim=-1)
    paired_text = F.normalize(paired_text.float(), dim=-1)
    similarity = particles @ paired_text.T
    target = similarity.diagonal()
    ranks = (similarity > target[:, None]).sum(-1) + 1
    return {
        "r@1": float((ranks <= 1).float().mean()),
        "r@5": float((ranks <= 5).float().mean()),
        "r@10": float((ranks <= 10).float().mean()),
        "median_rank": float(ranks.float().median()),
        "paired_cosine": float(target.mean()),
    }