from typing import Any, ClassVar, Literal
from typing_extensions import override

import jax
import jax.numpy as jnp
from jaxtyping import Array, Complex, Float

from ...coordinates import make_frequency_grid
from ...ndimage import (
    block_reduce_downsample,
    convert_fftn_to_rfftn,
    irfftn,
    operators as op,
    resize_with_crop_or_pad,
    rfftn,
)
from .._image_config import AbstractImageConfig
from .._volume import IndependentAtomVolume
from .base_integrator import AbstractVolumeIntegrator


try:
    import jax_finufft as jnufft

    JAX_FINUFFT_IMPORT_ERROR = None
except ModuleNotFoundError as err:
    jnufft = None
    JAX_FINUFFT_IMPORT_ERROR = err


class FFTAtomProjection(
    AbstractVolumeIntegrator[IndependentAtomVolume],
    strict=True,
):
    """Integrate atomic parametrization of a volume onto
    the exit plane using non-uniform FFTs plus convolution.
    """

    sampling_mode: Literal["average", "point"]
    upsample_factor: int | None
    shape: tuple[int, int] | None
    eps: float
    opts: Any

    is_projection_approximation: ClassVar[bool] = True

    def __init__(
        self,
        *,
        sampling_mode: Literal["average", "point"] = "average",
        upsample_factor: int | None = None,
        shape: tuple[int, int] | None = None,
        eps: float = 1e-6,
        opts: Any = None,
    ):
        """**Arguments:**

        - `sampling_mode`:
            If `'average'`, convolve with a box function to sample the
            projected volume at a pixel to be the average value of the
            underlying continuous function. If `'point'`, the volume at
            a pixel will be point sampled.
        - `upsample_factor`:
            If provided, first compute an upsampled version of the
            image at pixel size `image_config.pixel_size / upsample_factor`.
            Then, downsample with `cryojax.ndimage.block_reduce_downsample`
            to locally average to the correct pixel size. This is useful
            for reducing aliasing.
        - `shape`:
            If given, first compute the image at `shape`, then
            pad or crop to `image_config.padded_shape`.
        - `eps`:
            See [`jax-finufft`](https://github.com/flatironinstitute/jax-finufft)
            for documentation.
        - `opts`:
            A `jax_finufft.options.Opts` or `jax_finufft.options.NestedOpts`
            dataclass.
            See [`jax-finufft`](https://github.com/flatironinstitute/jax-finufft)
            for documentation.
        """
        if jnufft is None:
            raise RuntimeError(
                "Tried to use the `FFTAtomProjection` "
                "class, but `jax-finufft` is not installed. "
                "See https://github.com/flatironinstitute/jax-finufft "
                "for installation instructions."
            ) from JAX_FINUFFT_IMPORT_ERROR
        if sampling_mode not in ["average", "point"]:
            raise ValueError(
                "`sampling_mode` in `FFTAtomProjection` "
                "must be either 'average' for averaging within a "
                "pixel or 'point' for point sampling. Got "
                f"`sampling_mode = {sampling_mode}`."
            )
        self.sampling_mode = sampling_mode
        self.upsample_factor = upsample_factor
        self.shape = shape
        self.eps = eps
        self.opts = opts

    def __check_init__(self):
        if self.upsample_factor is not None:
            if self.upsample_factor % 2 == 0:
                raise ValueError(
                    f"Set `upsample_factor = {self.upsample_factor}` when instantiating "
                    "`FFTAtomProjection`, but only odd `upsample_factor` are supported."
                )

    @override
    def integrate(
        self,
        volume_representation: IndependentAtomVolume,
        image_config: AbstractImageConfig,
        outputs_real_space: bool = False,
    ) -> (
        Complex[
            Array,
            "{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}",
        ]
        | Float[Array, "{image_config.padded_y_dim} {image_config.padded_x_dim}"]
    ):
        """Compute a projection from scattering factors per atom type
        from the `IndependentAtomVolume`.

        **Arguments:**

        - `volume_representation`:
            The volume representation.
        - `image_config`:
            The configuration of the resulting image.
        - `outputs_real_space`:
            If `True`, return the image in real space. Otherwise,
            return in fourier.

        **Returns:**

        The integrated volume in real or fourier space at the
        `AbstractImageConfig.padded_shape`.
        """  # noqa: E501
        u = self.upsample_factor
        pixel_size = image_config.pixel_size
        shape = image_config.padded_shape if self.shape is None else self.shape
        if u is None:
            shape_u, pixel_size_u = shape, pixel_size
        else:
            shape_u, pixel_size_u = (u * shape[0], u * shape[1]), pixel_size / u
        if shape_u == image_config.padded_shape:
            frequency_grid = image_config.padded_full_frequency_grid_in_angstroms
        else:
            frequency_grid = make_frequency_grid(
                shape_u, pixel_size_u, outputs_rfftfreqs=False
            )
        frequency_grid = jnp.fft.fftshift(frequency_grid, axes=(0, 1))
        proj_kernel = lambda pos, kernel: _project_with_nufft(
            shape_u,
            pixel_size_u,
            pos,
            kernel,
            frequency_grid,
            eps=self.eps,
            opts=self.opts,
        )
        # Compute projection over atom types
        fourier_projection = jax.tree.reduce(
            lambda x, y: x + y,
            jax.tree.map(
                proj_kernel,
                volume_representation.position_pytree,
                volume_representation.scattering_factor_pytree,
                is_leaf=lambda x: isinstance(x, op.AbstractFourierOperator),
            ),
        )
        # Apply anti-aliasing filter
        if self.sampling_mode == "average":
            antialias_fn = op.FourierSinc(box_width=pixel_size_u)
            fourier_projection *= antialias_fn(frequency_grid)
        # Shift zero frequency component to corner and convert to
        # rfft
        fourier_projection = convert_fftn_to_rfftn(
            jnp.fft.ifftshift(fourier_projection), mode="real"
        )
        if self.shape is None:
            if u is None:
                return (
                    irfftn(fourier_projection, s=shape)
                    if outputs_real_space
                    else fourier_projection
                )
            else:
                projection = _block_average(irfftn(fourier_projection, s=shape_u), u)
                return projection if outputs_real_space else rfftn(projection)
        else:
            projection = irfftn(fourier_projection, s=shape_u)
            if u is not None:
                projection = _block_average(projection, u)
            projection = resize_with_crop_or_pad(projection, image_config.padded_shape)
            return projection if outputs_real_space else rfftn(projection)


def _project_with_nufft(shape, ps, pos, kernel, freqs, eps=1e-6, opts=None):
    assert jnufft is not None
    # Get x and y coordinates
    positions_xy = pos[:, :2]
    # Normalize coordinates betweeen -pi and pi
    ny, nx = shape
    box_xy = ps * jnp.asarray((nx, ny))
    positions_periodic = 2 * jnp.pi * positions_xy / box_xy
    # Unpack
    x, y = positions_periodic[:, 0], positions_periodic[:, 1]
    n_atoms = x.size
    area_element = ps**2
    # Compute
    fourier_projection = kernel(freqs) * (
        jnufft.nufft1(
            shape,
            jnp.full((n_atoms,), 1.0 + 0.0j),
            y,
            x,
            eps=eps,
            opts=opts,
            iflag=-1,
        )
        / area_element
    )

    return fourier_projection


def _block_average(x, factor):
    return block_reduce_downsample(x, factor, jax.lax.add) / factor**x.ndim
