"""Some of the code in this submodule is taken directly from the MIT-licensed
's2fft' package. Below is the original license text:

Copyright (c) 2022 Authors & Contributors

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

from functools import partial

import jax
import jax.numpy as jnp
from jax.scipy.spatial.transform import Rotation


def _compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
    """from s2fft.recursions.risbo_jax"""

    if el == 0:
        dl = dl.at[el + L - 1, el + L - 1].set(1.0)
        return dl
    if el == 1:
        cosb = jnp.cos(beta)
        sinb = jnp.sin(beta)

        coshb = jnp.cos(beta / 2.0)
        sinhb = jnp.sin(beta / 2.0)
        sqrt2 = jnp.sqrt(2.0)

        dl = dl.at[L - 2, L - 2].set(coshb**2)
        dl = dl.at[L - 2, L - 1].set(sinb / sqrt2)
        dl = dl.at[L - 2, L].set(sinhb**2)

        dl = dl.at[L - 1, L - 2].set(-sinb / sqrt2)
        dl = dl.at[L - 1, L - 1].set(cosb)
        dl = dl.at[L - 1, L].set(sinb / sqrt2)

        dl = dl.at[L, L - 2].set(sinhb**2)
        dl = dl.at[L, L - 1].set(-sinb / sqrt2)
        dl = dl.at[L, L].set(coshb**2)
        return dl
    else:
        coshb = -jnp.cos(beta / 2.0)
        sinhb = jnp.sin(beta / 2.0)
        dd = jnp.zeros((2 * el + 2, 2 * el + 2))

        # First pass
        j = 2 * el - 1
        i = jnp.arange(j)
        k = jnp.arange(j)

        sqrt_jmk = jnp.sqrt(j - k)
        sqrt_kp1 = jnp.sqrt(k + 1)
        sqrt_jmi = jnp.sqrt(j - i)
        sqrt_ip1 = jnp.sqrt(i + 1)

        dlj = dl[k - (el - 1) + L - 1][:, i - (el - 1) + L - 1]

        dd = dd.at[:j, :j].add(
            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_jmk, optimize=True) * dlj * coshb
        )
        dd = dd.at[:j, 1 : j + 1].add(
            jnp.einsum("i,k->ki", -sqrt_ip1, sqrt_jmk, optimize=True) * dlj * sinhb
        )
        dd = dd.at[1 : j + 1, :j].add(
            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_kp1, optimize=True) * dlj * sinhb
        )
        dd = dd.at[1 : j + 1, 1 : j + 1].add(
            jnp.einsum("i,k->ki", sqrt_ip1, sqrt_kp1, optimize=True) * dlj * coshb
        )

        dl = dl.at[-el + L - 1 : el + 1 + L - 1, -el + L - 1 : el + 1 + L - 1].multiply(
            0.0
        )

        j = 2 * el
        i = jnp.arange(j)
        k = jnp.arange(j)

        # Second pass
        sqrt_jmk = jnp.sqrt(j - k)
        sqrt_kp1 = jnp.sqrt(k + 1)
        sqrt_jmi = jnp.sqrt(j - i)
        sqrt_ip1 = jnp.sqrt(i + 1)

        dl = dl.at[-el + L - 1 : el + L - 1, -el + L - 1 : el + L - 1].add(
            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_jmk, optimize=True)
            * dd[:j, :j]
            * coshb,
        )
        dl = dl.at[-el + L - 1 : el + L - 1, L - el : L + el].add(
            jnp.einsum("i,k->ki", -sqrt_ip1, sqrt_jmk, optimize=True)
            * dd[:j, :j]
            * sinhb,
        )
        dl = dl.at[L - el : L + el, -el + L - 1 : el + L - 1].add(
            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_kp1, optimize=True)
            * dd[:j, :j]
            * sinhb,
        )
        dl = dl.at[L - el : L + el, L - el : L + el].add(
            jnp.einsum("i,k->ki", sqrt_ip1, sqrt_kp1, optimize=True)
            * dd[:j, :j]
            * coshb,
        )
        return dl / ((2 * el) * (2 * el - 1))


def _generate_rotate_dls(L: int, beta: float) -> jnp.ndarray:
    """*from s2fft.utils*"""
    dls = []
    dl_iter = jnp.zeros((2 * L - 1, 2 * L - 1)).astype(float)
    for el in range(L):
        dl_iter = _compute_full(dl_iter, beta, L, el)
        dls.append(dl_iter)
    return jnp.stack(dls, axis=0)


def compute_rotation_matrices(deg, x, y, z, theta, homogeneous=False):

    def kron(m, n):
        """Kronecker delta. Can we do better?"""
        return m == n

    def Umn(m, n):
        """Compute the (m, n) term of the transformation
        matrix from complex to real Ylms. This part is from starry"""
        if n < 0:
            term1 = 1j
        elif n == 0:
            term1 = jnp.sqrt(2) / 2
        else:
            term1 = 1
        if (m > 0) and (n < 0) and (n % 2 == 0):
            term2 = -1
        elif (m > 0) and (n > 0) and (n % 2 != 0):
            term2 = -1
        else:
            term2 = 1
        return term1 * term2 * 1 / jnp.sqrt(2) * (kron(m, n) + kron(m, -n) + 0j)

    def U(l):
        """Compute the complete U transformation matrix.. This part is from starry"""
        res = jnp.zeros((2 * l + 1, 2 * l + 1)) + 0j
        for m in range(-l, l + 1):
            for n in range(-l, l + 1):
                res = res.at[m + l, n + l].set(Umn(m, n))
        return res

    def euler(x, y, z, theta):
        """axis-angle to euler angles"""
        # the jnp where for theta == 0 is to avoid nans when computing grad
        axis = jnp.array([jnp.where(theta == 0.0, 1.0, x), y, z])
        _theta = jnp.where(theta == 0.0, 1.0, theta)
        axis = axis / jnp.linalg.norm(axis)
        r = Rotation.from_rotvec(axis * _theta)
        return jnp.where(theta == 0.0, jnp.array([0.0, 0.0, 0.0]), r.as_euler("zyz"))

    alpha, beta, gamma = euler(x, y, z, theta)
    dls = _generate_rotate_dls(deg + 1, beta)
    N = jnp.arange(-deg, deg + 1)
    m, mp = jnp.meshgrid(N, N)
    Dlm = jnp.exp(-1j * (m * alpha + mp * gamma)) * dls

    u = U(deg)
    u_inv = jnp.conj(u.T)

    # The output of generate_rotate_dls has shape (L, 2L+1, 2L+1)
    Rs_full = jnp.real(u_inv[None, :] @ Dlm @ u[None, :])

    if homogeneous:
        Rs = Rs_full
    else:
        # reformat to a list of matrices with different shape
        # (to match current implementation of rotation.dot_rotation_matrix)
        Rs = []

        for i in range(deg + 1):
            Rs.append(Rs_full[i, deg - i : deg + i + 1, deg - i : deg + i + 1])

    return Rs


def __exp_array(L: int, x: float) -> jnp.ndarray:
    """Private function to generate rotation arrays for alpha/gamma rotations"""
    return jnp.exp(-1j * jnp.arange(-L + 1, L) * x)


@partial(jax.jit, static_argnums=(1))
def rotate_flms(
    flm: jnp.ndarray,
    L: int,
    rotation: tuple[float, float, float],
    dl_array: jnp.ndarray = None,
) -> jnp.ndarray:
    """Rotates an array of spherical harmonic coefficients by angle rotation.

    Args:
        flm (jnp.ndarray): Array of spherical harmonic coefficients.
        L (int): Harmonic band-limit.
        rotation  (Tuple[float, float, float]): Rotation on the sphere
            (alpha, beta, gamma).
        dl_array (jnp.ndarray, optional): Precomputed array of reduced Wigner d-function
            coefficients, see :func:~`generate_rotate_dls`. Defaults to None.

    Returns:
        jnp.ndarray: Rotated spherical harmonic coefficients with shape [L,2L-1].
    """

    # Split out angles
    alpha = __exp_array(L, rotation[0])
    gamma = __exp_array(L, rotation[2])
    beta = rotation[1]

    # Create empty arrays
    flm_rotated = jnp.zeros_like(flm)

    dl = (
        dl_array
        if dl_array is not None
        else jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.float64)
    )

    # Perform rotation
    for el in range(L):
        if dl_array is None:
            dl = _compute_full(dl, beta, L, el)
        n_max = min(el, L - 1)

        m = jnp.arange(-el, el + 1)
        n = jnp.arange(-n_max, n_max + 1)

        flm_rotated = flm_rotated.at[el, L - 1 + m].add(
            jnp.einsum(
                "mn,n->m",
                jnp.einsum(
                    "mn,m->mn",
                    (
                        dl[m + L - 1][:, n + L - 1]
                        if dl_array is None
                        else dl[el, m + L - 1][:, n + L - 1]
                    ),
                    alpha[m + L - 1],
                    optimize=True,
                ),
                gamma[n + L - 1] * flm[el, n + L - 1],
            )
        )
    return flm_rotated
