from functools import partial

import equinox
import jax
import jax.numpy as jnp
from jaxtyping import Array, ArrayLike, Float, PRNGKeyArray

from furax import IndexOperator, SymmetricBandToeplitzOperator
from furax.obs._detectors import DetectorArray

default_fft_size = SymmetricBandToeplitzOperator._get_default_fft_size

__all__ = [
    'GapFillingOperator',
]


class GapFillingOperator(equinox.Module):
    """Class for filling masked time samples with a constrained noise realization.

    We follow the gap-filling algorithm described in https://doi.org/10.1103/PhysRevD.65.022003,
    section II.C, page 6. It assumes that the noise is piece wise stationary and has Gaussian
    statistics, described by the covariance matrix ``cov``.

    Example:
    Gap-filling a single-detector timestream

        >>> detectors = DetectorArray(jnp.array(0), jnp.array(0), jnp.array(1))
        >>> key = jax.random.key(0)
        >>> key, subkey = jax.random.split(key)
        >>> nsamples = 10
        >>> x = jax.random.normal(subkey, detectors.shape + (nsamples,))
        >>> in_structure = jax.ShapeDtypeStruct(x.shape, x.dtype)
        >>> mask = jnp.array([1, 1, 1, 0, 0, 1, 1, 1, 1, 1], dtype=bool)
        >>> mask_op = IndexOperator(jnp.where(mask), in_structure=in_structure)
        >>> cov = SymmetricBandToeplitzOperator(jnp.array([1.0]), in_structure)
        >>> gf = GapFillingOperator(cov, mask_op, detectors)
        >>> gap_filled_x = gf(key, x)
        >>> assert gap_filled_x.shape == x.shape
        >>> assert jnp.all(gap_filled_x[mask] == x[mask])

    Attributes:
        cov: A SymmetricBandToeplitzOperator representing the noise covariance matrix.
        mask_op: An IndexOperator for masking the gaps.
        detectors: A DetectorArray representing the detectors.
        rate: The sampling rate of the data.
    """

    cov: SymmetricBandToeplitzOperator
    mask_op: IndexOperator
    detectors: DetectorArray
    rate: float = 1.0

    def __call__(self, key: PRNGKeyArray, x: Float[Array, ' *shape']) -> Float[Array, ' *shape']:
        """Performs the gap-filling operation.

        Args:
            key: A PRNG key generated by the user.
            x: The vector to be filled.
        """
        real: Float[Array, ' ...'] = self._generate_realization_for(x, key)
        p, u = self.mask_op, self.mask_op.T  # pack, unpack operators
        incomplete_cov = p @ self.cov @ u
        op = self.cov @ u @ incomplete_cov.I @ p
        y: Float[Array, ' ...'] = real + op(x - real)
        # copy valid samples from original vector
        y = y.at[p.indices].set(p(x))
        return y

    @staticmethod
    def folded_psd(
        n_tt: Float[ArrayLike, ' _'], fft_size: int
    ) -> Float[Array, ' {fft_size // 2 + 1}']:
        """Returns the folded Power Spectral Density of a one-dimensional vector.

        Args:
            n_tt: The autocorrelation function of the vector.
            fft_size: The size of the FFT to use (at least twice the size of ``n_tt``).
        """
        n_tt = jnp.asarray(n_tt)
        kernel = GapFillingOperator._get_kernel(n_tt, fft_size)
        psd = jnp.abs(jnp.fft.rfft(kernel, n=fft_size))
        # zero out DC value
        psd = psd.at[0].set(0)
        return psd

    @staticmethod
    def _get_kernel(n_tt: Float[Array, ' _'], size: int) -> Float[Array, ' {size}']:
        lagmax = n_tt.size - 1
        padding_size = size - (2 * lagmax + 1)
        if padding_size < 0:
            msg = f'The maximum lag ({lagmax}) is too large for the required kernel size ({size}).'
            raise ValueError(msg)
        kernel = jnp.concatenate((n_tt, jnp.zeros(padding_size), n_tt[-1:0:-1]))
        return kernel

    def _generate_realization_for(
        self, x: Float[Array, ' *shape'], key: PRNGKeyArray
    ) -> Float[Array, ' *shape']:
        @partial(jnp.vectorize, signature='(n),(k),()->(n)')
        def func(x, n_tt, subkey):  # type: ignore[no-untyped-def]
            x_size = x.size
            fft_size = default_fft_size(x_size)
            npsd = fft_size // 2 + 1
            norm = self.rate * float(npsd - 1)

            # Get PSD values (size = fft_size // 2 + 1)
            psd = self.folded_psd(n_tt, fft_size)
            scale = jnp.sqrt(norm * psd)

            # Gaussian Re/Im random numbers
            rngdata = jax.random.normal(subkey, shape=(fft_size,))

            fdata = jnp.empty(npsd, dtype=complex)

            # Set DC and Nyquist frequency imaginary parts to zero
            fdata = fdata.at[0].set(rngdata[0] + 0.0j)
            fdata = fdata.at[-1].set(rngdata[npsd - 1] + 0.0j)

            # Repack the other values
            fdata = fdata.at[1:-1].set(rngdata[1 : npsd - 1] + 1j * rngdata[-1 : npsd - 1 : -1])

            # scale by PSD and inverse FFT
            tdata = jnp.fft.irfft(fdata * scale)

            # subtract the DC level for the samples we want
            offset = (fft_size - x_size) // 2
            xi = tdata[offset : offset + x_size]
            xi -= jnp.mean(xi)
            return xi

        subkeys = self.detectors.split_key(key)
        real: Float[Array, ' ...'] = func(x, self.cov.band_values, subkeys)
        return real
