import torch


def mapping_shtools_to_compact(lmax):
    """
    pyshtools uses an output format to represent spherical harmonic coefficients that is not memory efficient.
    This function creates a mapping to represent the coefficients differently.

    Our representation Y(l, m), shape(lmax+1, lmax+1) :
        Y(0, 0) Y(1, 1) Y(2, 2) ...
        Y(1,-1) Y(1, 0) Y(2, 1) ...
        Y(2,-2) Y(2,-1) Y(2, 0) ...
        ...     ...     ...

    Example :
        mapping = mapping_shtools_to_compact(lmax)
        x, _ = SHExpandLSQ(d, phi, theta, lmax, [1, -1])  # pyshtools
        y = torch.tensor(x)[mapping[..., 0], mapping[..., 1], mapping[..., 2]]  # Compact
        z = torch.zeros([2, lmax+1, lmax+1])
        z[mapping[..., 0], mapping[..., 1], mapping[..., 2]] = y  # Back to pyshtools
    """
    mapping = torch.zeros([lmax + 1, lmax + 1, 3], dtype=torch.long)
    mapping[..., 0] = torch.tril(torch.ones([lmax + 1, lmax + 1], dtype=torch.long)) - torch.eye(lmax + 1,
                                                                                                 dtype=torch.long)
    linspace = torch.linspace(0, lmax, lmax + 1, dtype=torch.long)
    mapping[..., 1] = torch.triu(linspace.view(1, -1).expand(lmax + 1, lmax + 1)) \
        + torch.tril(linspace.view(-1, 1).expand(lmax + 1, lmax + 1) - torch.diag(linspace))
    mapping[..., 2] = torch.abs(linspace.view(1, -1).expand(lmax + 1, lmax + 1) -
                                linspace.view(-1, 1).expand(lmax + 1, lmax + 1))
    return mapping