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


# ---------------------------------------------------------------------------
# SH basis evaluator  (low-level helper)
# ---------------------------------------------------------------------------


def _eval_sh_bases_fast(basis_dim: int, dirs: Tensor) -> Tensor:
    """Evaluate spherical harmonics bases at unit directions up to degree 4.

    Uses the efficient recurrence described in:
    "Efficient Spherical Harmonic Evaluation", Peter-Pike Sloan, JCGT 2013
    https://jcgt.org/published/0002/02/06/

    See reference C++ code in https://jcgt.org/published/0002/02/06/code.zip

    Args:
        basis_dim: Number of SH basis functions to evaluate. Must be a square
            number in [1, 25] (i.e. 1, 4, 9, 16, or 25).
        dirs: Unit direction vectors of shape (..., 3).

    Returns:
        Tensor of shape (..., basis_dim) containing the evaluated SH basis
        values for each direction.
    """
    result = torch.empty(
        (*dirs.shape[:-1], basis_dim), dtype=dirs.dtype, device=dirs.device
    )

    result[..., 0] = 0.2820947917738781

    if basis_dim <= 1:
        return result

    x, y, z = dirs.unbind(-1)

    fTmpA = -0.48860251190292
    result[..., 2] = -fTmpA * z
    result[..., 3] = fTmpA * x
    result[..., 1] = fTmpA * y
    if basis_dim <= 4:
        return result

    z2 = z * z
    fTmpB = -1.092548430592079 * z
    fTmpA = 0.5462742152960395
    fC1 = x * x - y * y
    fS1 = 2 * x * y
    result[..., 6] = 0.9461746957575601 * z2 - 0.3153915652525201
    result[..., 7] = fTmpB * x
    result[..., 5] = fTmpB * y
    result[..., 8] = fTmpA * fC1
    result[..., 4] = fTmpA * fS1

    if basis_dim <= 9:
        return result

    fTmpC = -2.285228997322329 * z2 + 0.4570457994644658
    fTmpB = 1.445305721320277 * z
    fTmpA = -0.5900435899266435
    fC2 = x * fC1 - y * fS1
    fS2 = x * fS1 + y * fC1
    result[..., 12] = z * (1.865881662950577 * z2 - 1.119528997770346)
    result[..., 13] = fTmpC * x
    result[..., 11] = fTmpC * y
    result[..., 14] = fTmpB * fC1
    result[..., 10] = fTmpB * fS1
    result[..., 15] = fTmpA * fC2
    result[..., 9] = fTmpA * fS2

    if basis_dim <= 16:
        return result

    fTmpD = z * (-4.683325804901025 * z2 + 2.007139630671868)
    fTmpC = 3.31161143515146 * z2 - 0.47308734787878
    fTmpB = -1.770130769779931 * z
    fTmpA = 0.6258357354491763
    fC3 = x * fC2 - y * fS2
    fS3 = x * fS2 + y * fC2
    result[..., 20] = 1.984313483298443 * z2 * (
        1.865881662950577 * z2 - 1.119528997770346
    ) + -1.006230589874905 * (0.9461746957575601 * z2 - 0.3153915652525201)
    result[..., 21] = fTmpD * x
    result[..., 19] = fTmpD * y
    result[..., 22] = fTmpC * fC1
    result[..., 18] = fTmpC * fS1
    result[..., 23] = fTmpB * fC2
    result[..., 17] = fTmpB * fS2
    result[..., 24] = fTmpA * fC3
    result[..., 16] = fTmpA * fS3
    return result


# ---------------------------------------------------------------------------
# Spherical harmonics — public API
# ---------------------------------------------------------------------------


def _spherical_harmonics(
    degrees_to_use: int,
    dirs: torch.Tensor,  # [..., 3]
    coeffs: torch.Tensor,  # [..., K, 3] or [..., K, 1]
) -> torch.Tensor:
    """Evaluate spherical harmonics color from SH coefficients and view directions.

    Pure-PyTorch implementation of `gsplat.cuda._wrapper.spherical_harmonics()`.
    Directions are L2-normalized before basis evaluation, so raw (unnormalized)
    directions may be passed safely.

    Args:
        degrees_to_use: Highest SH degree to include; only the first
            ``(degrees_to_use + 1) ** 2`` coefficient bands are used.
        dirs: View direction vectors of shape (..., 3). Need not be unit length.
        coeffs: SH coefficients of shape (..., K, 3) for RGB or (..., K, 1)
            for monochrome, where K >= ``(degrees_to_use + 1) ** 2``.

    Returns:
        Tensor of shape (..., 3) or (..., 1) containing the evaluated color
        values, matching the last dimension of ``coeffs``.
    """
    assert (degrees_to_use + 1) ** 2 <= coeffs.shape[-2], coeffs.shape
    batch_dims = dirs.shape[:-1]
    assert dirs.shape == batch_dims + (3,), dirs.shape
    assert (
        (len(coeffs.shape) == len(batch_dims) + 2)
        and coeffs.shape[:-2] == batch_dims
        and (coeffs.shape[-1] == 3 or coeffs.shape[-1] == 1)
    ), coeffs.shape
    dirs = F.normalize(dirs, p=2, dim=-1)
    num_bases = (degrees_to_use + 1) ** 2
    bases = torch.zeros_like(coeffs[..., 0])
    bases[..., :num_bases] = _eval_sh_bases_fast(num_bases, dirs)
    return (bases[..., None] * coeffs).sum(dim=-2)