"""Functions that generate various matrices for the FMM problem.

Copyright (c) Meta Platforms, Inc. and affiliates.
"""

import functools
from typing import Tuple

import jax.numpy as jnp

from fmmax import _fft, basis, utils


def omega_script_k_matrix_patterned(
    wavelength: jnp.ndarray,
    z_permittivity_matrix: jnp.ndarray,
    transverse_permeability_matrix: jnp.ndarray,
    transverse_wavevectors: jnp.ndarray,
) -> jnp.ndarray:
    """Compute the omega-script-k matrix of equation 26 from [2012 Liu]."""
    # The script-k matrix from equation 19 of [2012 Liu].
    script_k_matrix = script_k_matrix_patterned(
        z_permittivity_matrix, transverse_wavevectors
    )
    angular_frequency = utils.angular_frequency_for_wavelength(wavelength)
    angular_frequency_squared = angular_frequency[..., jnp.newaxis, jnp.newaxis] ** 2
    return angular_frequency_squared * transverse_permeability_matrix - script_k_matrix


# -----------------------------------------------------------------------------
# Functions to compute the k- and script-k matrices.
# -----------------------------------------------------------------------------


def script_k_matrix_uniform(
    permittivity: jnp.ndarray,
    transverse_wavevectors: jnp.ndarray,
) -> jnp.ndarray:
    """Returns the uniform-layer script-k matrix from eq. 19 of [2012 Liu]."""
    kx = transverse_wavevectors[..., 0]
    ky = transverse_wavevectors[..., 1]
    return jnp.block(
        [
            [
                utils.diag(ky / permittivity[..., jnp.newaxis] * ky),
                utils.diag(-ky / permittivity[..., jnp.newaxis] * kx),
            ],
            [
                utils.diag(-kx / permittivity[..., jnp.newaxis] * ky),
                utils.diag(kx / permittivity[..., jnp.newaxis] * kx),
            ],
        ]
    )


def script_k_matrix_patterned(
    z_permittivity_matrix: jnp.ndarray,
    transverse_wavevectors: jnp.ndarray,
) -> jnp.ndarray:
    """Returns the patterned-layer script-k matrix from eq. 19 of [2012 Liu]."""
    dtype = jnp.promote_types(z_permittivity_matrix, transverse_wavevectors)
    kx = transverse_wavevectors[..., 0].astype(dtype)
    ky = transverse_wavevectors[..., 1].astype(dtype)
    z_inv_kx = jnp.linalg.solve(z_permittivity_matrix.astype(dtype), utils.diag(kx))
    z_inv_ky = jnp.linalg.solve(z_permittivity_matrix.astype(dtype), utils.diag(ky))
    return jnp.block(
        [
            [ky[..., :, jnp.newaxis] * z_inv_ky, -ky[..., :, jnp.newaxis] * z_inv_kx],
            [-kx[..., :, jnp.newaxis] * z_inv_ky, kx[..., :, jnp.newaxis] * z_inv_kx],
        ]
    )


def k_matrix_patterned(
    z_permeability_matrix: jnp.ndarray,
    transverse_wavevectors: jnp.ndarray,
) -> jnp.ndarray:
    """Returns the k-matrix for patterned magnetic materials."""
    dtype = jnp.promote_types(z_permeability_matrix, transverse_wavevectors)
    kx = transverse_wavevectors[..., 0].astype(dtype)
    ky = transverse_wavevectors[..., 1].astype(dtype)
    z_inv_kx = jnp.linalg.solve(z_permeability_matrix.astype(dtype), utils.diag(kx))
    z_inv_ky = jnp.linalg.solve(z_permeability_matrix.astype(dtype), utils.diag(ky))
    return jnp.block(
        [
            [kx[..., :, jnp.newaxis] * z_inv_kx, kx[..., :, jnp.newaxis] * z_inv_ky],
            [ky[..., :, jnp.newaxis] * z_inv_kx, ky[..., :, jnp.newaxis] * z_inv_ky],
        ]
    )


# -----------------------------------------------------------------------------
# Functions for computing Fourier convolution matrices for isotropic media.
# -----------------------------------------------------------------------------


def transverse_permittivity_fft(
    permittivity: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Computes the `eps` matrix from [2012 Liu] equation 15 using `fft` scheme.

    Args:
        permittivity: The permittivity array, with shape `(..., nx, ny)`.
        expansion: The field expansion to be used.

    Returns:
        The transverse permittivity matrix.
    """
    eps_hat = _fft.fourier_convolution_matrix(permittivity, expansion)
    zeros = jnp.zeros_like(eps_hat)
    return jnp.block([[eps_hat, zeros], [zeros, eps_hat]])


def transverse_permittivity_vector(
    permittivity: jnp.ndarray,
    tx: jnp.ndarray,
    ty: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Computes transverse permittivity matrix using a vector field methods.

    The transverse permittivity matrix is E2 given in equation 51 of [2012 Liu].

    Args:
        permittivity: The permittivity array, with shape `(..., nx, ny)`.
        tx: The x-component of the tangent vector field.
        ty: The y-component of the tangent vector field.
        expansion: The field expansion to be used.

    Returns:
        The `eps` matrix.
    """
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)

    eps_hat = _transform(permittivity)
    zeros = jnp.zeros_like(eps_hat)
    eps_matrix = jnp.block([[eps_hat, zeros], [zeros, eps_hat]])

    eta = 1 / permittivity
    eta_hat = _transform(eta)
    delta_hat = eps_hat - jnp.linalg.inv(eta_hat)
    delta_matrix = jnp.block([[delta_hat, zeros], [zeros, delta_hat]])

    Pxx, Pxy, Pyx, Pyy = _tangent_terms(tx, ty)
    p_matrix = jnp.block(
        [[_transform(Pyy), _transform(Pyx)], [_transform(Pxy), _transform(Pxx)]]
    )
    return eps_matrix - delta_matrix @ p_matrix


def _tangent_terms(
    tx: jnp.ndarray,
    ty: jnp.ndarray,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Returns the `(Pxx, Pxy, Pyx, Pyy)` from [2012 Liu] equation 50.

    Args:
        tx: The x-component of the tangent vector field.
        ty: The y-component of the tangent vector field.

    Returns:
        The terms `(Pxx, Pxy, Pyx, Pyy)`.
    """
    denom = jnp.abs(tx) ** 2 + jnp.abs(ty) ** 2
    denom_safe = jnp.where(jnp.isclose(denom, 0), 1.0, denom)

    # Compute terms from equation 50 of [2012 Liu]. Note that there is an error
    # in equation 48-51 of [2012 Liu], where x- and y- subscripts are swapped for
    # the diagonal blocks. This can be shown by deriving equation 48 from 46, and
    # can also easily be seen by considering the case of a 1D gratng periodic in
    # the x-direction. In this case, the tangent vector field has a y-component
    # only (i.e. tx = 0), and we want to recover the transverse permittivity
    # given by equation 43. This is possible only by swapping the x- and y-
    # subscripts in the diagonal terms of equations 48-51.
    #
    # Note that we also investigated both orderings numerically with the 1D
    # grating example, and found improved convergence with this correction.
    Pyy = jnp.abs(tx) ** 2 / denom_safe
    Pyx = jnp.conj(tx) * ty / denom_safe
    Pxy = tx * jnp.conj(ty) / denom_safe
    Pxx = jnp.abs(ty) ** 2 / denom_safe
    return Pxx, Pxy, Pyx, Pyy


# -----------------------------------------------------------------------------
# Functions for computing Fourier convolution matrices for anisotropic media.
# -----------------------------------------------------------------------------


def transverse_permittivity_fft_anisotropic(
    permittivity_xx: jnp.ndarray,
    permittivity_xy: jnp.ndarray,
    permittivity_yx: jnp.ndarray,
    permittivity_yy: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Compute anisotropic transverse permittivity matrix using the fft scheme."""
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)
    return jnp.block(
        [
            [_transform(permittivity_yy), _transform(-permittivity_yx)],
            [_transform(-permittivity_xy), _transform(permittivity_xx)],
        ]
    )


def transverse_permeability_fft_anisotropic(
    permeability_xx: jnp.ndarray,
    permeability_xy: jnp.ndarray,
    permeability_yx: jnp.ndarray,
    permeability_yy: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Compute anisotropic transverse permeability matrix using the fft scheme."""
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)
    return jnp.block(
        [
            [_transform(permeability_xx), _transform(permeability_xy)],
            [_transform(permeability_yx), _transform(permeability_yy)],
        ]
    )


def transverse_permittivity_vector_anisotropic(
    permittivity_xx: jnp.ndarray,
    permittivity_xy: jnp.ndarray,
    permittivity_yx: jnp.ndarray,
    permittivity_yy: jnp.ndarray,
    tx: jnp.ndarray,
    ty: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Compute the transverse permittivity matrix with a vector scheme.

    The transverse permittivity matrix E relates the electric and electric displacement
    fields, such that

        [-Dy, Dx]^T = E [-Ey, Ex]^T

    Args:
        permittivity_xx: The xx-component of the permittivity tensor, with
            shape `(..., nx, ny)`.
        permittivity_xy: The xy-component of the permittivity tensor.
        permittivity_yx: The yx-component of the permittivity tensor.
        permittivity_yy: The yy-component of the permittivity tensor.
        tx: The x-component of the tangent vector field.
        ty: The y-component of the tangent vector field.
        expansion: The field expansion to be used.

    Returns:
        The transverse permittivity matrix.
    """
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)

    # Define the real-space and Fourier transformed rotation matrices. The rotation
    # matrix is defined such that T [Et, En]^T = [-Ey, Ex].
    (
        rotation_matrix,
        fourier_rotation_matrix,
        fourier_inverse_rotation_matrix,
    ) = _rotation_matrices(
        t00=ty,
        t01=tx.conj(),
        t10=-tx,
        t11=ty.conj(),
        expansion=expansion,
    )

    # Obtain the permittivity tensor for rotated coordinates, i.e. coordinates where
    # unit vectors are tangent and normal to the vector field defined by `(tx, ty)`.
    permittivity_tensor = jnp.block(
        [
            [
                permittivity_yy[..., jnp.newaxis, jnp.newaxis],
                -permittivity_yx[..., jnp.newaxis, jnp.newaxis],
            ],
            [
                -permittivity_xy[..., jnp.newaxis, jnp.newaxis],
                permittivity_xx[..., jnp.newaxis, jnp.newaxis],
            ],
        ]
    )
    assert rotation_matrix.shape[-2:] == permittivity_tensor.shape[-2:] == (2, 2)
    rotated_permittivity_tensor = jnp.linalg.solve(
        rotation_matrix, permittivity_tensor @ rotation_matrix
    )

    # The Fourier permittivity matrix is the central matrix in equation 45 of
    # [2012 Liu]. The top left block relates the tangential E and D fields, while
    # the bottom right block relates the normal E and D fields. Consequently, these
    # use Laurent's rule and the inverse rule, respectively. Note that the
    # non-diagonal blocks also use Laurent's rule, as these tensor elements may be
    # zero, making the inverse rule problematic.
    fourier_permittivity_matrix = jnp.block(
        [
            [
                _transform(rotated_permittivity_tensor[..., 0, 0]),
                _transform(rotated_permittivity_tensor[..., 0, 1]),
            ],
            [
                _transform(rotated_permittivity_tensor[..., 1, 0]),
                jnp.linalg.inv(_transform(1 / rotated_permittivity_tensor[..., 1, 1])),
            ],
        ]
    )
    return (
        fourier_rotation_matrix
        @ fourier_permittivity_matrix
        @ fourier_inverse_rotation_matrix
    )


def transverse_permeability_vector_anisotropic(
    permeability_xx: jnp.ndarray,
    permeability_xy: jnp.ndarray,
    permeability_yx: jnp.ndarray,
    permeability_yy: jnp.ndarray,
    tx: jnp.ndarray,
    ty: jnp.ndarray,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Compute the transverse permeability matrix with a vector scheme.

    The transverse permeability matrix M relates the magnetic and magnetic flux
    density fields, such that

        [Bx, Bx]^T = M [Hx, Hy]^T

    Important differences from the `_transverse_permittivity_vector_anisotropic`
    function result from the different definitions of E and M matrices.

    Args:
        permeability_xx: The xx-component of the permeability tensor, with
            shape `(..., nx, ny)`.
        permeability_xy: The xy-component of the permeability tensor.
        permeability_yx: The yx-component of the permeability tensor.
        permeability_yy: The yy-component of the permeability tensor.
        tx: The x-component of the tangent vector field.
        ty: The y-component of the tangent vector field.
        expansion: The field expansion to be used.

    Returns:
        The transverse permeability matrix.
    """
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)

    # Define the real-space and Fourier transformed rotation matrices. The rotation
    # matrix is defined such that T [Ht, Hn]^T = [Hx, Hy], and differs from that
    # used for the transverse permittivity calculation.
    (
        rotation_matrix,
        fourier_rotation_matrix,
        fourier_inverse_rotation_matrix,
    ) = _rotation_matrices(
        t00=tx,
        t01=-ty.conj(),
        t10=ty,
        t11=tx.conj(),
        expansion=expansion,
    )

    # Obtain the permeability tensor for rotated coordinates, i.e. coordinates where
    # unit vectors are tangent and normal to the vector field defined by `(tx, ty)`.
    permeability_tensor = jnp.block(
        [
            [
                permeability_xx[..., jnp.newaxis, jnp.newaxis],
                permeability_xy[..., jnp.newaxis, jnp.newaxis],
            ],
            [
                permeability_yx[..., jnp.newaxis, jnp.newaxis],
                permeability_yy[..., jnp.newaxis, jnp.newaxis],
            ],
        ]
    )
    assert rotation_matrix.shape[-2:] == permeability_tensor.shape[-2:] == (2, 2)
    rotated_permeability_tensor = jnp.linalg.solve(
        rotation_matrix, permeability_tensor @ rotation_matrix
    )

    # The Fourier permeability matrix is analogous to the central matrix in equation 45
    # of [2012 Liu], in that its top-left block relates the tangential magnetic and
    # magnetic flux density fields, and the bottom-right block relates the normal
    # fields. Consequently, these use Laurent's rule and the inverse rule, respectively.
    # Note that the non-diagonal blocks also use Laurent's rule, as these tensor
    # elements may be zero, making the inverse rule problematic.
    fourier_permeability_matrix = jnp.block(
        [
            [
                _transform(rotated_permeability_tensor[..., 0, 0]),  # tt
                _transform(rotated_permeability_tensor[..., 0, 1]),  # tn
            ],
            [
                _transform(rotated_permeability_tensor[..., 1, 0]),  # nt
                jnp.linalg.inv(
                    _transform(1 / rotated_permeability_tensor[..., 1, 1])
                ),  # nn
            ],
        ]
    )
    return (
        fourier_rotation_matrix
        @ fourier_permeability_matrix
        @ fourier_inverse_rotation_matrix
    )


def _rotation_matrices(
    t00: jnp.ndarray,
    t01: jnp.ndarray,
    t10: jnp.ndarray,
    t11: jnp.ndarray,
    expansion: basis.Expansion,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Generate real-space and Fourier transformed rotation matrices."""
    rotation_matrix = jnp.block(
        [
            [
                t00[..., jnp.newaxis, jnp.newaxis],
                t01[..., jnp.newaxis, jnp.newaxis],
            ],
            [
                t10[..., jnp.newaxis, jnp.newaxis],
                t11[..., jnp.newaxis, jnp.newaxis],
            ],
        ]
    )

    # Blockwise Fourier transform of the rotation matrix.
    _transform = functools.partial(_fft.fourier_convolution_matrix, expansion=expansion)
    fourier_rotation_matrix = jnp.block(
        [
            [_transform(t00), _transform(t01)],
            [_transform(t10), _transform(t11)],
        ]
    )

    # Blockwise Fourier transform of the inverse rotation matrix.
    denom = jnp.abs(t00 * t11 - t10 * t01)
    denom_safe = jnp.where(jnp.isclose(denom, 0), 1.0, denom)
    fourier_inverse_rotation_matrix = jnp.block(
        [
            [_transform(t11 / denom_safe), _transform(-t01 / denom_safe)],
            [_transform(-t10 / denom_safe), _transform(t00 / denom_safe)],
        ]
    )
    return rotation_matrix, fourier_rotation_matrix, fourier_inverse_rotation_matrix
