"""Defines functions related to farfield patterns in the FMM scheme.

These functions are intended for calculations that involve Brillouin
zone integration, and batched in-plane wavevectors as generated by
`basis.brillouin_zone_in_plane_wavevector`.

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

from typing import Callable, Tuple

import jax
import jax.numpy as jnp
import numpy as onp

from fmmax import basis, misc, utils


def farfield_profile(
    flux: jnp.ndarray,
    wavelength: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Computes a farfield profile.

    This function effectively "unstacks" the values for each Fourier order and
    for each point in the Brillouin zone sampling scheme.

    Args:
        flux: The flux array, with shape ``(..., num_bz_kx, num_bz_ky, ...,
            2 * num_terms, num_sources)``.
        wavelength: The wavelength, batch-compatible with `flux`.
        in_plane_wavevector: The in-plane wavevector for the zeroth Fourier
            order, batch-compatible with `flux`.
        primitive_lattice_vectors: The primitive lattice vectors of the unit cell.
        expansion: The expansion used for the fields.
        brillouin_grid_axes: Specifies the two axes of ``flux`` corresponding to
            the Brillouin zone grid.

    Returns:
        The polar and azimuthal angles, solid angle associated with each value,
        and the farfield power.
    """
    assert flux.shape[-2] == 2 * expansion.num_terms
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, flux.ndim
    )

    ndim_batch = flux.ndim - 2
    wavelength = misc.atleast_nd(wavelength, ndim_batch)
    in_plane_wavevector = misc.atleast_nd(in_plane_wavevector, ndim_batch + 1)

    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector,
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
    )
    transverse_wavevectors = _unflatten_transverse_wavevectors(
        transverse_wavevectors, expansion, brillouin_grid_axes
    )
    flux = _unflatten_flux(flux, expansion, brillouin_grid_axes)

    # Remove the brillouin zone axes from `wavelength`, making it compatible
    # with the unflattened flux and transverse wavevectors.
    wavelength = jnp.squeeze(wavelength, axis=brillouin_grid_axes)

    polar_angle, azimuthal_angle = _angles_from_unflattened_transverse_wavevectors(
        transverse_wavevectors=transverse_wavevectors,
        wavelength=wavelength,
    )

    # Transform flux form units of power per unit Brillouin zone area to
    # power per unit solid angle. Add dummy dimensions for the polarization
    # and sources.
    solid_angle = _solid_angle_from_unflattened_transverse_wavevectors(
        transverse_wavevectors=transverse_wavevectors,
        wavelength=wavelength,
    )
    transformed_flux = flux / solid_angle[..., jnp.newaxis, jnp.newaxis]
    return polar_angle, azimuthal_angle, solid_angle, transformed_flux


def _angles_from_unflattened_transverse_wavevectors(
    transverse_wavevectors: jnp.ndarray,
    wavelength: jnp.ndarray,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Computes the propagation angles in free space for given wavevectors.

    Evanescent modes whose transverse wavevector magnitude exceeds that of
    the free space wavevector are given a polar angle of ``pi / 2``.

    Args:
        transverse_wavevectors: The unflattened transverse wavectors, with
            shape ``(..., nkx, nky, 2)``.
        wavelength: The free-space wavelength.

    Returns:
        Arrays containing the polar and azimuthal angles.
    """
    assert transverse_wavevectors.ndim - 3 == wavelength.ndim
    kx = transverse_wavevectors[..., 0]
    ky = transverse_wavevectors[..., 1]
    kt = jnp.sqrt(kx**2 + ky**2)

    sin_polar_angle = kt * wavelength[..., jnp.newaxis, jnp.newaxis] / (2 * jnp.pi)
    polar_angle = jnp.where(
        sin_polar_angle > 1, jnp.pi / 2, jnp.arcsin(sin_polar_angle)
    )

    azimuthal_angle = jnp.angle(kx + 1j * ky)
    return polar_angle, azimuthal_angle


def _solid_angle_from_unflattened_transverse_wavevectors(
    transverse_wavevectors: jnp.ndarray,
    wavelength: jnp.ndarray,
) -> jnp.ndarray:
    """Computes the solid angle associated with each transverse wavevector.

    The transverse wavevectors should be unflattened, i.e. the ``-3`` and ``-2``
    axes should correspond to different points in k-space.

    Args:
        transverse_wavevectors: The unflattened transverse wavevectors
        wavelength: The free-space wavelength.

    Returns:
        The solid angle, with the shape matching the leading dimensions of
        `transverse_wavevectors`.
    """
    assert transverse_wavevectors.ndim >= 3

    kx = transverse_wavevectors[..., 0]
    ky = transverse_wavevectors[..., 1]

    # Normalize the wavevectors so they lie on the unit sphere.
    kx /= 2 * jnp.pi / wavelength[..., jnp.newaxis, jnp.newaxis]
    ky /= 2 * jnp.pi / wavelength[..., jnp.newaxis, jnp.newaxis]
    polar_angle = jnp.arcsin(jnp.sqrt(kx**2 + ky**2))

    # Each of our transverse wavevectors lies within a "cell" in the kxky plane.
    # Compute the area of each cell, and then project it onto the unit sphere
    # to get the solid angle associated with each transverse wavevector.
    #
    # First, compute the locations of the cell verteces.
    pad_width = ((0, 0),) * (kx.ndim - 2) + ((1, 1), (1, 1))
    vertex_kx = jnp.pad(kx, pad_width, mode="edge")
    vertex_ky = jnp.pad(ky, pad_width, mode="edge")
    vertex_kx = (
        vertex_kx[..., :-1, :-1]
        + vertex_kx[..., :-1, 1:]
        + vertex_kx[..., 1:, :-1]
        + vertex_kx[..., 1:, 1:]
    ) / 4
    vertex_ky = (
        vertex_ky[..., :-1, :-1]
        + vertex_ky[..., :-1, 1:]
        + vertex_ky[..., 1:, :-1]
        + vertex_ky[..., 1:, 1:]
    ) / 4
    vertex_kt = jnp.stack([vertex_kx, vertex_ky], axis=-1)

    # Find the vectors defining each parallelogramic cell.
    v1 = vertex_kt[..., :-1, 1:, :] - vertex_kt[..., :-1, :-1, :]
    v2 = vertex_kt[..., 1:, :-1, :] - vertex_kt[..., :-1, :-1, :]

    # Find the area of each parallelogramic cell.
    cell_area = jnp.abs(v1[..., 0] * v2[..., 1] - v2[..., 0] * v1[..., 1])

    # Project the area onto the unit sphere.
    projected_area = cell_area / jnp.cos(polar_angle)
    return projected_area


# -----------------------------------------------------------------------------
# Functions for computing the total flux in some angular cone.
# -----------------------------------------------------------------------------


def farfield_integrated_flux(
    flux: jnp.ndarray,
    wavelength: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
    angle_bounds_fn: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
    upsample_factor: int,
) -> jnp.ndarray:
    """Computes the flux within the bounds defined by ``angle_bounds_fn``.

    Args:
        flux: The flux array, with shape ``(..., num_bz_kx, num_bz_ky, ...
            2 * num_terms, num_sources)``.
        wavelength: The wavelength, batch-compatible with ``flux``.
        in_plane_wavevector: The in-plane wavevector for the zeroth Fourier
            order, batch-compatible with ``flux``.
        primitive_lattice_vectors: The primitive lattice vectors of the unit cell.
        expansion: The expansion used for the fields.
        brillouin_grid_axes: Specifies the two axes of ``flux`` corresponding to
            the Brillouin zone grid.
        angle_bounds_fn: A function with signature ``fn(polar_angle,
            azimuthal_angle) -> mask`` returning a mask that is ``True`` for angles
            that should be included in the integral.
        upsample_factor: Integer factor specifying upsampling performed in the
            integral, which is used to approximate trapezoidal rule integration.

    Returns:
        The integrated flux, with shape equal to the batch dimensions of flux,
        excluding those for the brillouin zone grid.
    """
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, flux.ndim
    )

    # Compute the weights array, which reduce the integration weights to
    # an inner product.
    weights = _integrated_flux_weights(
        flux=flux,
        wavelength=wavelength,
        in_plane_wavevector=in_plane_wavevector,
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
        brillouin_grid_axes=brillouin_grid_axes,
        angle_bounds_fn=angle_bounds_fn,
        upsample_factor=upsample_factor,
    )

    # Sum over the Brillouin zone and Fourier order axes.
    return jnp.sum(weights * flux, axis=brillouin_grid_axes + (-2,))


def _integrated_flux_weights(
    flux: jnp.ndarray,
    wavelength: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
    angle_bounds_fn: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
    upsample_factor: int,
) -> jnp.ndarray:
    """Returns the integration weights for the bounds defined by ``angle_bounds_fn``."""
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, flux.ndim
    )

    def _integrated_fn(flux):
        assert flux.shape[-1] == 1
        return jnp.sum(
            _integrated_flux_upsampled(
                flux=flux,
                wavelength=wavelength,
                in_plane_wavevector=in_plane_wavevector,
                primitive_lattice_vectors=primitive_lattice_vectors,
                expansion=expansion,
                brillouin_grid_axes=brillouin_grid_axes,
                angle_bounds_fn=angle_bounds_fn,
                upsample_factor=upsample_factor,
            )
        )

    # The weights are just the gradient of the integrated flux with respect
    # to the flat array elements; use a dummy flux having a single source.
    dummy_flux = jnp.ones_like(flux, shape=flux.shape[:-1] + (1,))
    return jax.grad(_integrated_fn)(dummy_flux)


def _integrated_flux_upsampled(
    flux: jnp.ndarray,
    wavelength: jnp.ndarray,
    in_plane_wavevector: jnp.ndarray,
    primitive_lattice_vectors: basis.LatticeVectors,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
    angle_bounds_fn: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray],
    upsample_factor: int,
) -> jnp.ndarray:
    """Computes the flux within the bounds defined by ``angle_bounds_fn``."""
    assert upsample_factor >= 1
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, flux.ndim
    )

    ndim_batch = flux.ndim - 2
    wavelength = misc.atleast_nd(wavelength, ndim_batch)
    in_plane_wavevector = misc.atleast_nd(in_plane_wavevector, ndim_batch + 1)

    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector=in_plane_wavevector,
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
    )

    flux = _unflatten_flux(flux, expansion, brillouin_grid_axes)
    transverse_wavevectors = _unflatten_transverse_wavevectors(
        transverse_wavevectors, expansion, brillouin_grid_axes
    )

    # Remove the `nan`s that are found at array locations having no associated
    # value in the original flattened arrays.
    flux = jnp.where(jnp.isnan(flux), 0, flux)
    transverse_wavevectors = jnp.where(
        jnp.isnan(transverse_wavevectors), 0, transverse_wavevectors
    )

    # Upsample the transverse wavevectors and flux to high resolution.
    assert flux.shape[-2] == 2
    upsampled_flux_shape = flux.shape[:-4] + (
        upsample_factor * flux.shape[-4],  # kx axis
        upsample_factor * flux.shape[-3],  # ky axis
        flux.shape[-2],  # Polarization axis, length 2
        flux.shape[-1],  # Source axis
    )
    flux = jax.image.resize(flux, upsampled_flux_shape, method="linear")

    assert transverse_wavevectors.shape[-1] == 2
    upsampled_wavevector_shape = transverse_wavevectors.shape[:-3] + (
        upsample_factor * transverse_wavevectors.shape[-3],  # kx axis
        upsample_factor * transverse_wavevectors.shape[-2],  # ky axis
        transverse_wavevectors.shape[-1],  # Direction axis.
    )
    transverse_wavevectors = jax.image.resize(
        transverse_wavevectors, upsampled_wavevector_shape, method="linear"
    )

    # Remove the brillouin grid axes from wavelength, and insert axes
    # for the unstacked wavevector.
    wavelength = jnp.squeeze(wavelength, brillouin_grid_axes)

    # Compute polar and azimuthal angles.
    polar_angle, azimuthal_angle = _angles_from_unflattened_transverse_wavevectors(
        transverse_wavevectors, wavelength
    )

    selected = angle_bounds_fn(polar_angle, azimuthal_angle)
    selected = jnp.where(jnp.isnan(selected), False, selected)
    selected = selected[..., jnp.newaxis, jnp.newaxis]

    masked_flux = jnp.where(selected, flux, 0)

    # Sum over the kx, ky, and polarization axes.
    return jnp.sum(masked_flux, axis=(-4, -3, -2)) / upsample_factor**2


# -----------------------------------------------------------------------------
# Functions related to unflattening results of Brillouin zone integration.
# -----------------------------------------------------------------------------


def _unflatten(flat: jnp.ndarray, expansion: basis.Expansion) -> jnp.ndarray:
    """Unflattens an array for a given expansion and Brillouin integration scheme.

    The returned array combines the values associated with all terms in the
    Fourier expansion at all points in the Brillouin zone grid in a single
    array with trailing axes havving shape ``(num_kx, num_ky)``. Elements in the
    output which have no corresponding elements in ``flat`` are given a value
    of ``nan``.

    The flat array should have shape ``(..., num_bz_kx, num_bz_ky, num_terms)``,
    where ``num_terms`` is the number of terms in the Fourier expansion, and the
    ``-3`` and ``-2`` axes are for the Brillouin zone grid, as used e.g. with
    Brillouin zone integration to model localized sources.

    This function assumes that the Brillouin zone is sampled on a regular grid,
    as produced by ``basis.brillouin_zone_in_plane_wavevector``.

    Args:
        flat: The flat array, with shape  ``(..., num_bz_kx, num_bz_ky, num_terms)``.
        expansion: The expansion used for the array.

    Returns:
        The unflattened array, with shape ``(batch_shape, num_kx, num_ky)``.
    """
    assert flat.ndim >= 3
    assert flat.shape[-1] == expansion.num_terms

    i = expansion.basis_coefficients[:, 0]
    j = expansion.basis_coefficients[:, 1]

    batch_shape = flat.shape[:-3]
    bz_grid_shape = flat.shape[-3:-1]

    # The shape of the output array shoudl accomodate all `(i, j)` values.
    shape = batch_shape + (
        (max(i) - min(i) + 1) * bz_grid_shape[0],
        (max(j) - min(j) + 1) * bz_grid_shape[1],
    )

    bz_i, bz_j = onp.meshgrid(
        onp.arange(bz_grid_shape[0]),
        onp.arange(bz_grid_shape[1]),
        indexing="ij",
    )

    offset_i = min(i) * bz_grid_shape[0]
    offset_j = min(j) * bz_grid_shape[1]

    merged_i = i * bz_grid_shape[0] + bz_i[..., onp.newaxis] - offset_i
    merged_j = j * bz_grid_shape[1] + bz_j[..., onp.newaxis] - offset_j

    stacked_i = merged_i.flatten()
    stacked_j = merged_j.flatten()
    stacked_flat = jnp.reshape(flat, batch_shape + (-1,))

    unflattened = jnp.full(shape, jnp.nan, dtype=flat.dtype)
    unflattened = unflattened.at[..., stacked_i, stacked_j].set(stacked_flat)
    return unflattened


def _unflatten_flux(
    flux: jnp.ndarray,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
) -> jnp.ndarray:
    """Unflattens a flux for a given expansion and Brillouin integration scheme.

    Args:
        flux: The flux array, with shape ``(..., num_bz_kx, num_bz_ky, ...
            2 * num_terms, num_sources)``.
        expansion: The expansion used for the flux.
        brillouin_grid_axes: The axes associated with the Brillouin zone grid.

    Returns:
        The unflattened flux, with shape ``(..., num_kx, num_ky, 2, num_sources)``.
    """
    assert flux.ndim >= 4
    assert flux.shape[-2] == 2 * expansion.num_terms
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, flux.ndim
    )

    # The flux array has values for two polarizations at each Fourier order. Split
    # these and treat them as a batch dimension.
    flux = jnp.reshape(flux, flux.shape[:-2] + (2, -1, flux.shape[-1]))

    # Transpose so the axes associated with the Fourier orders and Brillouin zone
    # grid are the trailing axes, as needed by `unflatten`.
    batch_axes = tuple(
        [i for i in range(flux.ndim) if i not in brillouin_grid_axes + (flux.ndim - 2,)]
    )
    axes = batch_axes + tuple([i for i in range(flux.ndim) if i not in batch_axes])
    flux = jnp.transpose(flux, axes)

    flux = _unflatten(flux, expansion)

    # Transpose so that polarization and sources are returned to the trailing axes.
    axes = tuple(range(flux.ndim - 4)) + (-2, -1, -4, -3)
    return jnp.transpose(flux, axes)


def _unflatten_transverse_wavevectors(
    transverse_wavevectors: jnp.ndarray,
    expansion: basis.Expansion,
    brillouin_grid_axes: Tuple[int, int],
) -> jnp.ndarray:
    """Unflatten transverse wavevectors for given expansion and integration scheme.

    Args:
        transverse_wavevectors: The transverse wavevectors array, with shape
            `(..., num_bz_kx, num_bz_ky, ..., num_terms, 2)`.
        expansion: The expansion used for the flux.
        brillouin_grid_axes: The axes associated with the Brillouin zone grid.

    Returns:
        The unflattened wavevectors, with shape `(..., num_kx, num_ky, 2)`.
    """
    assert transverse_wavevectors.ndim >= 4
    assert transverse_wavevectors.shape[-2:] == (expansion.num_terms, 2)
    brillouin_grid_axes = utils.absolute_axes(  # type: ignore[assignment]
        brillouin_grid_axes, transverse_wavevectors.ndim
    )

    # Transpose so the axes associated with the Fourier orders and Brillouin zone
    # grid are the trailing axes, as needed by `unflatten`.
    non_batch_axes = brillouin_grid_axes + (transverse_wavevectors.ndim - 2,)
    batch_axes = tuple(
        [i for i in range(transverse_wavevectors.ndim) if i not in non_batch_axes]
    )
    axes = batch_axes + tuple(
        [i for i in range(transverse_wavevectors.ndim) if i not in batch_axes]
    )
    transverse_wavevectors = jnp.transpose(transverse_wavevectors, axes)

    transverse_wavevectors = _unflatten(transverse_wavevectors, expansion)

    # Transpose so the trailing axis is for the wavevector direction.
    axes = tuple(range(transverse_wavevectors.ndim - 3)) + (-2, -1, -3)
    return jnp.transpose(transverse_wavevectors, axes)
