"""Functions related to sources in the FMM algorithm.

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

from typing import Optional, Tuple

import jax.numpy as jnp

from fmmax import basis, fft, fields, fmm, scattering, utils


def amplitudes_for_fields(
    ex: jnp.ndarray,
    ey: jnp.ndarray,
    hx: jnp.ndarray,
    hy: jnp.ndarray,
    layer_solve_result: fmm.LayerSolveResult,
    brillouin_grid_axes: Optional[Tuple[int, int]] = None,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Computes the amplitudes for fields.

    The fields are expected to be given for a supercell containing one or more unit
    cells, i.e. for a calculation with in-plane wavevector generated by
    `basis.brillouin_zone_in_plane_wavevector`. The shape of the supercell is
    inferred from the batch shape of the `layer_solve_result` and the specified
    `brillouin_grid_axes`.

    Args:
        ex: The x-oriented electric field on the real-space grid, with a trailing
            batch dimension.
        ey: The y-oriented electric field on the real-space grid, with shape
            `(..., nx, ny, num_fields)`.
        hx: The x-oriented magnetic field on the real-space grid.
        hy: The y-oriented magnetic field on the real-space grid.
        layer_solve_result: The eigensolve result for the layer containing the fields.
        brillouin_grid_axes: The batch axes corresponding to the Brillouin zone grid.

    Returns:
        The forward and backward wave amplitudes.
    """

    if not (ex.shape == ey.shape == hx.shape == hy.shape) or (ex.ndim < 3):
        raise ValueError(
            f"All fields must have rank of at least 3 with matching shape, but got "
            f"shapes of {ex.shape}, {ey.shape}, {hx.shape}, and {hy.shape}."
        )

    if not utils.batch_compatible_shapes(
        layer_solve_result.batch_shape,
        ex.shape[:-3],
        ey.shape[:-3],
        hx.shape[:-3],
        hy.shape[:-3],
    ):
        raise ValueError(
            f"Fields must be batch-compatible with `layer_solve_result`, but got "
            f"shapes of {ex.shape}, {ey.shape}, {hx.shape}, and {hy.shape} with "
            f"`layer_solve_result` batch shape {layer_solve_result.batch_shape}."
        )

    if brillouin_grid_axes is None:
        brillouin_grid_shape = (1, 1)
    else:
        absolute_brillouin_grid_axes = utils.absolute_axes(brillouin_grid_axes, ex.ndim)
        brillouin_grid_shape = (
            layer_solve_result.batch_shape[absolute_brillouin_grid_axes[0]],
            layer_solve_result.batch_shape[absolute_brillouin_grid_axes[1]],
        )

    if (
        ex.shape[-3] % brillouin_grid_shape[0] != 0
        or ex.shape[-2] % brillouin_grid_shape[1] != 0
    ):
        raise ValueError(
            f"Field shapes must be evenly divisible by the Brillouin grid shape, but "
            f"got shape {ex.shape} when inferred Brillouin grid shape is "
            f"{brillouin_grid_shape}."
        )

    unit_cell_shape = (
        ex.shape[-3] // brillouin_grid_shape[0],
        ex.shape[-2] // brillouin_grid_shape[1],
    )

    x, y = basis.unit_cell_coordinates(
        primitive_lattice_vectors=layer_solve_result.primitive_lattice_vectors,
        shape=unit_cell_shape,
        num_unit_cells=brillouin_grid_shape,
    )
    kx = layer_solve_result.in_plane_wavevector[..., 0, jnp.newaxis, jnp.newaxis]
    ky = layer_solve_result.in_plane_wavevector[..., 1, jnp.newaxis, jnp.newaxis]
    phase = jnp.exp(1j * (kx * x + ky * y))[..., jnp.newaxis]

    def _process_field(field: jnp.ndarray) -> jnp.ndarray:
        field /= phase
        field_split = jnp.asarray(
            tuple(
                jnp.split(s, indices_or_sections=brillouin_grid_shape[1], axis=-2)
                for s in jnp.split(
                    field, indices_or_sections=brillouin_grid_shape[0], axis=-3
                )
            )
        )
        field_fft = fft.fft(
            field_split, expansion=layer_solve_result.expansion, axes=(-3, -2)
        )
        return jnp.sum(field_fft, axis=(0, 1))

    ex = _process_field(ex)
    ey = _process_field(ey)
    hx = _process_field(hx)
    hy = _process_field(hy)

    stacked_fields = jnp.concatenate([-ey, ex, hx, hy], axis=-2)
    matrix = fields.field_conversion_matrix(layer_solve_result)
    amplitudes = jnp.linalg.solve(matrix, stacked_fields)

    fwd_amplitude, bwd_amplitude = jnp.split(amplitudes, indices_or_sections=2, axis=-2)
    return fwd_amplitude, bwd_amplitude


# -----------------------------------------------------------------------------
# Functions related to internal sources.
# -----------------------------------------------------------------------------


def dirac_delta_source(
    location: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Returns the coefficients for a delta source at the specified location.

    This function is appropriate for creating sources to be used with
    `amplitudes_for_source`.

    Args:
        location: The location of the source, with shape `(num_sources, 2)` and
            the trailing axis giving the x and y location. By convention, the
            center of the unit cell is at `(0, 0)`.
        in_plane_wavevector: The in-plane wavevevector for the calculation, which
            gives the offset of the plane wave decomposition. Has shape `(..., 2)`
            with possible batch dimensions.
        primitive_lattice_vectors: The primitive lattice vectors of the unit cell.
        expansion: The Fourier expansion used for the calculation.

    Returns:
        The coefficients, with the shape `(..., expansion.num_terms, num_sources)`.
    """
    if location.ndim != 2 or location.shape[-1] != 2:
        raise ValueError(
            f"`location` must be rank-2 with a trailing axis size of 2, but got "
            f"shape {location.shape}"
        )
    if in_plane_wavevector.shape[-1] != 2:
        raise ValueError(
            f"`in_plane_wavevector` must have a trailing axis size of 2, but got "
            f"shape {in_plane_wavevector.shape}"
        )

    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector=in_plane_wavevector,
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
    )
    kx = transverse_wavevectors[..., 0][..., jnp.newaxis]
    ky = transverse_wavevectors[..., 1][..., jnp.newaxis]
    x = location[:, 0]
    y = location[:, 1]
    return jnp.exp(-1j * (kx * x + ky * y))


def gaussian_source(
    fwhm: jnp.ndarray,
    location: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
) -> jnp.ndarray:
    """Returns the coefficients for a Gaussian source at the specified location.

    This function is appropriate for creating sources to be used with
    `amplitudes_for_source`.

    Args:
        fwhm: The full-width at half-maximum for the Gaussian source.
        location: The location of the source, with shape `(num_sources, 2)` and
            the trailing axis giving the x and y location. By convention, the
            center of the unit cell is at `(0, 0)`.
        in_plane_wavevector: The in-plane wavevevector for the calculation, which
            gives the offset of the plane wave decomposition. Has shape `(..., 2)`
            with possible batch dimensions.
        primitive_lattice_vectors: The primitive lattice vectors of the unit cell.
        expansion: The Fourier expansion used for the calculation.

    Returns:
        The coefficients, with the shape `(..., expansion.num_terms, num_sources)`.
    """
    if location.ndim != 2 or location.shape[-1] != 2:
        raise ValueError(
            f"`location` must be rank-2 with a trailing axis size of 2, but got "
            f"shape {location.shape}"
        )
    if in_plane_wavevector.shape[-1] != 2:
        raise ValueError(
            f"`in_plane_wavevector` must have a trailing axis size of 2, but got "
            f"shape {in_plane_wavevector.shape}"
        )

    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector=in_plane_wavevector,
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
    )
    kx = transverse_wavevectors[..., 0][..., jnp.newaxis]
    ky = transverse_wavevectors[..., 1][..., jnp.newaxis]
    x = location[:, 0]
    y = location[:, 1]
    coeffs = jnp.exp(-1j * (kx * x + ky * y))

    sigma = fwhm / (2 * jnp.sqrt(2 * jnp.log(2)))
    return coeffs * jnp.exp(-(sigma**2) * (kx**2 + ky**2) / 2)


def amplitudes_for_source(
    jx: jnp.ndarray,
    jy: jnp.ndarray,
    jz: jnp.ndarray,
    s_matrix_before_source: scattering.ScatteringMatrix,
    s_matrix_after_source: scattering.ScatteringMatrix,
) -> Tuple[
    jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray
]:
    """Computes wave amplitudes resulting from an internal source.

    The configuration of the calculation is depicted below. A source is located at
    the interface of two layers. The layer stacks before and after the source are
    named as such. The function computes the amplitudes of forward-going and
    backward-going waves at various locations within the stack, also depicted below.

                                source _____
                                            |
             __before_______________________V___after_____________________

            |             |     |           |           |     |           |
            |   layer 0   | ... |  layer i  |   layer   | ... |  layer N  |
            | start layer |     |           |   i + 1   |     | end layer |
            |             |     |           |           |     |           |

                                 -> a_i      -> a_i+1          -> a_N
                    b_0 <-            b_i <-    b_i+1 <-

    Args:
        jx: The x-oriented dipole amplitude; must be at least rank-3 with a
            trailing batch dimension.
        jy: The y-oriented dipole amplitude, with shape matching `jx`.
        jz: The z-oriented dipole amplitude, with shape matching `jx`.
        s_matrix_before_source: The scattering matrix for the layer substack
            before the source, having no overlap with the after-source substack.
            Scattering matrix pairs returned by `scattering.stack_s_matrices_interior`
            may not be directly used.
        s_matrix_after_source: The scattering matrix for the layer substack after
            the source.

    Returns:
        The wave amplitudes:
            backward_amplitude_0_end: The backward-going wave amplitude at the end
                of the first layer.
            forward_amplitude_before_start: The forward-going wave amplitude at the
                start of the layer preceding the source.
            backward_amplitude_before_end: The backward-going wave amplitude at the
                end of the layer preceding the source, i.e. just before the source.
            forward_amplitude_after_start: The forward-going wave amplitude at the
                start of the layer following the source, i.e. just after the source.
            backward_amplitude_after_end: The backward-going wave amplitude at the
                end of the layer following the source.
            forward_amplitude_N_start: The forward-going wave amplitude at the start
                of the final layer.
    """
    pol = polarization_terms(
        jx, jy, jz, layer_solve_result=s_matrix_before_source.end_layer_solve_result
    )
    matrix = emission_matrix(
        s_matrix_before_source=s_matrix_before_source,
        s_matrix_after_source=s_matrix_after_source,
    )

    # The matrix solve directly yields the forward-going wave amplitudes after the
    # source, and the backward-going wave amplitudes before the source.
    ab = jnp.linalg.solve(matrix, pol)
    fwd_amplitude_after_start, bwd_amplitude_before_end = jnp.split(ab, 2, axis=-2)

    # Other wave amplitudes are computed from the scattering matrices.
    bwd_amplitude_0_end = s_matrix_before_source.s22 @ bwd_amplitude_before_end
    fwd_amplitude_before_start = s_matrix_before_source.s12 @ bwd_amplitude_before_end
    bwd_amplitude_after_end = s_matrix_after_source.s21 @ fwd_amplitude_after_start
    fwd_amplitude_N_start = s_matrix_after_source.s11 @ fwd_amplitude_after_start

    return (
        bwd_amplitude_0_end,
        fwd_amplitude_before_start,
        bwd_amplitude_before_end,
        fwd_amplitude_after_start,
        bwd_amplitude_after_end,
        fwd_amplitude_N_start,
    )


def polarization_terms(
    jx: jnp.ndarray,
    jy: jnp.ndarray,
    jz: jnp.ndarray,
    layer_solve_result: fmm.LayerSolveResult,
) -> jnp.ndarray:
    """Computes the polarization terms for currents on the real-space grid.

    The polarization terms are discussed in section 7 of [1999 Whittaker].

    Args:
        jx: The Fourier-transformed x-oriented dipole amplitude, with a
            trailing batch dimension.
        jy: The Fourier-transformed y-oriented dipole amplitude.
        jz: The Fourier-transformed z-oriented dipole amplitude.
        layer_solve_result: The results of the layer eigensolve.

    Returns:
        The polarization vector containing parallel and z-oriented terms.
    """
    if not (jx.shape == jy.shape == jz.shape) or jx.ndim < 2:
        raise ValueError(
            f"`jx`, `jy`, and `jz` must have matching shape and be at least "
            f"rank-2, but got shapes {jx.shape, jy.shape, jz.shape}."
        )

    # Add dimensions as needed to ensures shapes are compatible.
    eta_matrix = layer_solve_result.inverse_z_permittivity_matrix
    leading_ndim = max(jx.ndim - 3, eta_matrix.ndim - 2)
    eta_matrix = utils.atleast_nd(eta_matrix, n=(leading_ndim + 2))
    jx = utils.atleast_nd(jx, n=(leading_ndim + 2))
    jy = utils.atleast_nd(jy, n=(leading_ndim + 2))
    jz = utils.atleast_nd(jz, n=(leading_ndim + 2))

    p_parallel = jnp.concatenate([jy, -jx], axis=-2)

    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector=layer_solve_result.in_plane_wavevector,
        primitive_lattice_vectors=layer_solve_result.primitive_lattice_vectors,
        expansion=layer_solve_result.expansion,
    )
    kx = transverse_wavevectors[..., 0, jnp.newaxis]
    ky = transverse_wavevectors[..., 1, jnp.newaxis]

    # Use the `eta_matrix`, the Fourier convolution matrix for the inverse of
    # the permittivity rather than inverting the Fourier convolution matrix for
    # permittivity itself. The improved convergence due to this choice is
    # discussed in [1999 Whittaker].
    pzy = ky * eta_matrix @ jz
    pzx = kx * eta_matrix @ jz
    pz = jnp.concatenate([-pzy, pzx], axis=-2)

    p_parallel = jnp.broadcast_to(p_parallel, pz.shape)
    return jnp.concatenate([pz, p_parallel], axis=-2)


def emission_matrix(
    s_matrix_before_source: scattering.ScatteringMatrix,
    s_matrix_after_source: scattering.ScatteringMatrix,
) -> jnp.ndarray:
    """Returns the emission matrix for a source between two layers.

    Args:
        s_matrix_before_source: The scattering matrix for the layer substack
            before the source, having no overlap with the after-source substack.
            Scattering matrix pairs returned by `scattering.stack_s_matrices_interior`
            may not be directly used.
        s_matrix_after_source: The scattering matrix for the layer substack after
            the source.

    Returns:
        The emission matrix.
    """
    q = s_matrix_before_source.end_layer_solve_result.eigenvalues
    phi = s_matrix_before_source.end_layer_solve_result.eigenvectors
    omega_k = s_matrix_before_source.end_layer_solve_result.omega_script_k_matrix
    layer_thickness = s_matrix_before_source.end_layer_thickness

    next_q = s_matrix_after_source.start_layer_solve_result.eigenvalues
    next_phi = s_matrix_after_source.start_layer_solve_result.eigenvectors
    next_omega_k = s_matrix_after_source.start_layer_solve_result.omega_script_k_matrix
    next_layer_thickness = s_matrix_after_source.start_layer_thickness

    # Compute terms used in the matrix and assemble the matrix.
    fd = jnp.exp(1j * q * layer_thickness)
    next_fd = jnp.exp(1j * next_q * next_layer_thickness)

    fd_s12_before = utils.diag(fd) @ s_matrix_before_source.s12
    fd_s21_after = utils.diag(next_fd) @ s_matrix_after_source.s21

    shape = jnp.broadcast_shapes(q.shape, next_q.shape)
    eye = utils.diag(jnp.ones(shape, dtype=q.dtype))

    # Equation 7.9 from [1999 Whittaker].
    matrix = jnp.block(
        [
            [
                next_omega_k @ next_phi @ utils.diag(1 / next_q) @ (eye - fd_s21_after),
                omega_k @ phi @ utils.diag(1 / q) @ (eye - fd_s12_before),
            ],
            [
                next_phi @ (eye + fd_s21_after),
                -phi @ (eye + fd_s12_before),
            ],
        ]
    )

    return matrix
