from __future__ import annotations

import torch


def to_simplex_prob(v: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    x = v.abs() + float(eps)
    s = x.sum(dim=-1, keepdim=True)
    s = torch.where(s > 0, s, torch.ones_like(s))
    p = x / s
    return p.clamp_min(float(eps))


def fisher_rao_distance_simplex(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    p_s = to_simplex_prob(p, eps=eps)
    q_s = to_simplex_prob(q, eps=eps)
    inner = torch.sum(torch.sqrt(p_s * q_s + float(eps)), dim=-1)
    inner = inner.clamp(0.0, 1.0)
    return 2.0 * torch.acos(inner.clamp(0.0, 1.0))


def bures_distance_simplex(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    p_s = to_simplex_prob(p, eps=eps)
    q_s = to_simplex_prob(q, eps=eps)
    inner = torch.sum(torch.sqrt(p_s * q_s + float(eps)), dim=-1).clamp(0.0, 1.0)
    distance_sq = 2.0 * (1.0 - inner)
    return torch.sqrt(distance_sq.clamp_min(float(eps)))


def geodesic_interpolate_simplex(
    p: torch.Tensor,
    q: torch.Tensor,
    t: float = 0.5,
    eps: float = 1e-12,
) -> torch.Tensor:
    p_s = to_simplex_prob(p, eps=eps)
    q_s = to_simplex_prob(q, eps=eps)

    sp = torch.sqrt(p_s)
    sq = torch.sqrt(q_s)

    dot = torch.sum(sp * sq, dim=-1, keepdim=True).clamp(0.0, 1.0)
    omega = torch.acos(dot)
    sin_omega = torch.sin(omega)

    t_t = torch.as_tensor(t, dtype=sp.dtype, device=sp.device)
    t_t = t_t.reshape([1] * (sp.dim() - 1) + [1])

    small = sin_omega.abs() < 1e-10
    a = torch.where(small, 1.0 - t_t, torch.sin((1.0 - t_t) * omega) / sin_omega)
    b = torch.where(small, t_t, torch.sin(t_t * omega) / sin_omega)

    s_interp = a * sp + b * sq
    p_interp = s_interp * s_interp
    return to_simplex_prob(p_interp, eps=eps)


def geodesic_mean_simplex(
    basins: torch.Tensor,
    weights: torch.Tensor | None = None,
    eps: float = 1e-12,
) -> torch.Tensor:
    if basins.dim() != 2:
        raise ValueError("basins must have shape (N, D)")

    p = to_simplex_prob(basins, eps=eps)
    s = torch.sqrt(p)

    if weights is None:
        w = torch.full((p.size(0),), 1.0 / float(p.size(0)), dtype=p.dtype, device=p.device)
    else:
        if weights.dim() != 1 or int(weights.numel()) != int(p.size(0)):
            raise ValueError("weights must have shape (N,)")
        w = weights.to(dtype=p.dtype, device=p.device)
        w = w / (w.sum() + float(eps))

    mean_s = torch.sum(s * w.unsqueeze(-1), dim=0)
    norm = torch.linalg.vector_norm(mean_s).clamp_min(float(eps))
    mean_s = mean_s / norm

    mean_p = mean_s * mean_s
    return to_simplex_prob(mean_p, eps=eps)


__all__ = [
    "to_simplex_prob",
    "fisher_rao_distance_simplex",
    "bures_distance_simplex",
    "geodesic_interpolate_simplex",
    "geodesic_mean_simplex",
]
