from keras.src.backend.config import backend
from keras.src.backend import any_symbolic_tensors
from keras.src.ops.operation import Operation
from keras import ops
from keras import KerasTensor
from typing import Union, Tuple

if backend() == 'jax':
    from .jax import fft_fn, ifft_fn, fft2_fn, ifft2_fn, fft3_fn, ifft3_fn, rfft_fn, irfft_fn, rfft2_fn, irfft2_fn, rfft3_fn, irfft3_fn
    from .jax import fft_derivative as _fft_derivative

if backend() == 'tensorflow':
    from .tensorflow import fft_fn, ifft_fn, fft2_fn, fft3_fn, ifft3_fn, ifft2_fn, rfft_fn, irfft_fn, rfft2_fn, irfft2_fn, rfft3_fn, irfft3_fn
    from .tensorflow import fft_derivative as _fft_derivative

def cast_to_complex(x: Union[Tuple[KerasTensor, KerasTensor], KerasTensor]) -> Tuple[KerasTensor, KerasTensor]:
    if isinstance(x, tuple):
        return x
    else:
        return x, ops.zeros_like(x, dtype=x.dtype)


# === FFT ===
class FFT(Operation):
    """
    Base 1-D fast Fourier transform (FFT) operation.
    Keras-backend-agnostic version of the FFT.

    Notes
    -----
    Keras3 does not support complex dtypes.
    Therefore, the `keras.ops.fft` is quite cumbersume to use, since it explicitly requires the definition of a real and a complex part when calling `fft` ::

    >>> from keras import ops
    >>> x = ops.ones((4,))
    >>> ops.fft(x)
    ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
    >>> ops.fft((x, ops.zeros_like(x)))
    array([4.,  0.,  0.,  0.], dtype=float32), array([0., 0., 0., 0.], dtype=float32)
    
    This implementation attempts to increase the user-friendliness of the FFT call in Keras by handling the input automatically.

    """

    def __init__(self):
        super().__init__()
        self.fft_fn = fft_fn

    def compute_output_spec(self, x):
        """
        Compute output spec of Fourier transform

        Parameters
        ----------
        x : KerasTensor | tuple | list
            Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

        Returns
        -------
        y_real_spec, y_imag_spec : (KerasTensor, KerasTensor)
            spec of real- and imaginary part of `FFT(x)`

        """

        if not isinstance(x, (tuple, list)) or len(x) != 2:
            real = x
            imag = ops.zeros_like(x)
        else:
            real, imag = x
        # Both real and imaginary parts should have the same shape.
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )

        return (
            KerasTensor(shape=real.shape, dtype=real.dtype),
            KerasTensor(shape=imag.shape, dtype=imag.dtype),
        )

    def call(self, x):
        """
        Call method of FFT

        Parameters
        ----------
        x : KerasTensor | tuple | list
            Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

        Returns
        -------
        y_real, y_imag : (KerasTensor, KerasTensor)
            real- and imaginary part of `FFT(x)`

        Notes
        -----
        The input can be either a KerasTensor or a tuple or list of length 2, containing the real- and imaginary part of `x`.

        """

        return self.fft_fn(x)
    

class FFT2(FFT):
    """
    2-D fast Fourier transform

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = fft2_fn


class FFT3(FFT):
    """
    3-D fast Fourier transform

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = fft3_fn


class RFFT(FFT):
    """
    1-D fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = rfft_fn


class RFFT2(FFT):
    """
    2-D fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = rfft2_fn


class RFFT3(FFT):
    """
    3-D fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = rfft3_fn


# === IFFT ===
class IFFT(FFT):
    """
    1-D inverse fast Fourier transform

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = ifft_fn


class IFFT2(FFT):
    """
    2-D inverse fast Fourier transform

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = ifft2_fn


class IFFT3(FFT):
    """
    3-D inverse fast Fourier transform

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = ifft3_fn


class IRFFT(FFT):
    """
    1-D inverse fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = irfft_fn

    def compute_output_spec(self, x):
        """
        Compute output spec of Fourier transform

        Parameters
        ----------
        x : KerasTensor | tuple | list
            Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

        Returns
        -------
        y_real_spec : KerasTensor
            spec of real part of `IRFFT(x)`

        """

        if not isinstance(x, (tuple, list)) or len(x) != 2:
            real = x
            imag = ops.zeros_like(x)
        else:
            real, imag = x
        # Both real and imaginary parts should have the same shape.
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )

        return KerasTensor(shape=real.shape, dtype=real.dtype)


class IRFFT2(IRFFT):
    """
    2-D inverse fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = irfft2_fn


class IRFFT3(IRFFT):
    """
    3-D inverse fast Fourier transform for real-valued inputs

    """

    def __init__(self):
        super.__init__()
        self.fft_fn = irfft3_fn


# === now make function wrapper ===
# === forward FFT ===
def fft(x):
    """
    1-D fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of FFT(x).

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
    >>> y_real, y_imag = fft(signal)
    >>> ops.convert_to_numpy(y_real)
    array([ 25.       ,   3.3639607, -10.       ,  -9.36396  ,  -9.       ,
        -9.36396  , -10.       ,   3.3639607], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([ 0.       , -7.9497476,  1.       , -1.9497476,  0.       ,
        1.9497476, -1.       ,  7.9497476], dtype=float32)

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return FFT().symbolic_call(x)
    return fft_fn(x)


def fft2(x):
    """
    2-D fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of FFT2(x).

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.reshape(ops.array([-2, 8, 6, 12], dtype=float), (2, 2))
    >>> y_real, y_imag = fft2(signal)
    >>> ops.convert_to_numpy(y_real)
    array([
        [ 16.,  -8.],
        [ -4., -12.]
    ], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([
        [0., 0.],
        [0., 0.]
    ], dtype=float32)

    """
    
    if any_symbolic_tensors(cast_to_complex(x)):
        return FFT2().symbolic_call(x)
    return fft2_fn(x)


def fft3(x):
    """
    3-D fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of FFT3(x).

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.reshape(ops.array([-2, 8, 6, 12], dtype=float), (2, 2))
    >>> y_real, y_imag = fft3(signal)
    >>> ops.convert_to_numpy(y_real)
    array([
        [ 16.,  -8.],
        [ -4., -12.]
    ], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([
        [0., 0.],
        [0., 0.]
    ], dtype=float32)

    """
    
    if any_symbolic_tensors(cast_to_complex(x)):
        return FFT3().symbolic_call(x)
    return fft3_fn(x)


def rfft(x):
    """
    1-D fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor
        Real-valued input to FFT.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of RFFT(x).
        Both have size (n // 1 + 1) given ops.shape(x) == n

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
    >>> y_real, y_imag = rfft(signal)
    >>> ops.convert_to_numpy(y_real)
    array([ 25.       ,   3.3639607, -10.       ,  -9.36396  ,  -9.       ], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([ 0.       , -7.9497476,  1.       , -1.9497476,  0.       ], dtype=float32)
       
    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return RFFT().symbolic_call(x)
    return rfft_fn(x)


def rfft2(x):
    """
    2-D fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor
        Real-valued input to FFT.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of RFFT2(x).
        Both have size (n, n // 1 + 1) given ops.shape(x) == (n, n)

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.reshape(ops.array([-2, 8, 6, 12, 4, 0, 2, 2, -2, 8, 6, 12, 4, 0, 2, 2], dtype=float), (4, 4))
    >>> y_real, y_imag = rfft2(signal)
    >>> ops.convert_to_numpy(y_real)
    array([
        [ 64., -12., -24.],
        [  0.,   0.,   0.],
        [ 32., -20., -40.],
        [  0.,   0.,   0.]
    ], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([
        [ 0., 12.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  4.,  0.],
        [ 0.,  0.,  0.]
    ], dtype=float32)

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return RFFT2().symbolic_call(x)
    return rfft2_fn(x)


def rfft3(x):
    """
    3-D fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor
        Real-valued input to FFT.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of RFFT3(x).
        Both have size (n, n // 1 + 1) given ops.shape(x) == (n, n)

    Examples
    --------
    >>> from keras import ops
    >>> signal = ops.reshape(ops.array([-2, 8, 6, 12, 4, 0, 2, 2, -2, 8, 6, 12, 4, 0, 2, 2], dtype=float), (4, 4))
    >>> y_real, y_imag = rfft3(signal)
    >>> ops.convert_to_numpy(y_real)
    array([
        [ 64., -12., -24.],
        [  0.,   0.,   0.],
        [ 32., -20., -40.],
        [  0.,   0.,   0.]
    ], dtype=float32)
    >>> ops.convert_to_numpy(y_imag)
    array([
        [ 0., 12.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  4.,  0.],
        [ 0.,  0.,  0.]
    ], dtype=float32)

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return RFFT3().symbolic_call(x)
    return rfft3_fn(x)


# === inverse FFT ===
def ifft(x):
    """
    1-D inverse fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IFFT(x).

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return IFFT().symbolic_call(x)
    return ifft_fn(x)


def ifft2(x):
    """
    2-D inverse fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IFFT2(x).

    """
    
    if any_symbolic_tensors(cast_to_complex(x)):
        return IFFT2().symbolic_call(x)
    return ifft2_fn(x)


def ifft3(x):
    """
    3-D inverse fast Fourier transform

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IFFT3(x).

    """
    
    if any_symbolic_tensors(cast_to_complex(x)):
        return IFFT3().symbolic_call(x)
    return ifft3_fn(x)


def irfft(x, n=None):
    """
    1-D inverse fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.
    n : int, optional
        Length of the transformed axis of the output. For `n` output points, `n//2+1` input points are necessary.
        If the input is longer than tis, it is cropped.
        If it is shorter than this, it is padded with zeros.
        If `n` is not given, it is taken to be `2`(m-1)` where `m` is the length of the input along the axis specified by axis,
        cf. https://numpy.org/doc/stable/reference/generated/numpy.fft.irfft.html

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IRFFT(x).
        Both have size (n // 1 + 1) given ops.shape(x) == n

    """
    
    if any_symbolic_tensors(cast_to_complex(x)):
        return IRFFT().symbolic_call(x, n=n)
    return irfft_fn(x, n=n)


def irfft2(x, n=None):
    """
    2-D inverse fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.
    n : tuple, optional
        Shape (length of each transformed axis) of the output.
        `s` is also the number of input points used along this axis, except for the last axis,
        where `s[-1]//2+1` points of the input are used.
        Along any axis, if the shape indicated by `s` is smaller than that of the input, the input is cropped.
        If it is larger, the input is padded with zeros,
        cf. https://numpy.org/doc/stable/reference/generated/numpy.fft.irfftn.html
        
    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IRFFT2(x).
        Both have size (n // 1 + 1, n) given ops.shape(x) == (n, n)

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return IRFFT2().symbolic_call(x, n=n)
    return irfft2_fn(x, n=n)


def irfft3(x, n=None):
    """
    3-D inverse fast Fourier transform for real-valued inputs

    Parameters
    ----------
    x : KerasTensor | tuple | list
        Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.
    n : tuple, optional
        Shape (length of each transformed axis) of the output.
        `s` is also the number of input points used along this axis, except for the last axis,
        where `s[-1]//2+1` points of the input are used.
        Along any axis, if the shape indicated by `s` is smaller than that of the input, the input is cropped.
        If it is larger, the input is padded with zeros,
        cf. https://numpy.org/doc/stable/reference/generated/numpy.fft.irfftn.html

    Returns
    -------
    y_real, y_imag : (KerasTensor, KerasTensor)
        Tuple of real- and imaginary part of IRFFT3(x).
        Both have size (n // 1 + 1, n) given ops.shape(x) == (n, n)

    """

    if any_symbolic_tensors(cast_to_complex(x)):
        return IRFFT3().symbolic_call(x, n=n)
    return irfft3_fn(x, n=n)


# === FFT Derivative ===
class FFTDerivative(Operation):
    def __init__(self):
        super().__init__()

    def compute_output_spec(self, x):
        """
        Compute output spec of Fourier transform

        Parameters
        ----------
        x : KerasTensor | tuple | list
            Real- or complex input to FFT. A complex input must be composed of a tuple or list of the real- and imaginary part `(x_real, x_imag)`.

        Returns
        -------
        diff_spec : KerasTensor
            spec of derivative of x

        """
        if not isinstance(x, (tuple, list)) or len(x) != 2:
            real = x
            imag = ops.zeros_like(x)
        else:
            real, imag = x
        # Both real and imaginary parts should have the same shape.
        if real.shape != imag.shape:
            raise ValueError(
                "Input `x` should be a tuple of two tensors - real and "
                "imaginary. Both the real and imaginary parts should have the "
                f"same shape. Received: x[0].shape = {real.shape}, "
                f"x[1].shape = {imag.shape}"
            )
        return KerasTensor(shape=real.shape, dtype=real.dtype)
    
    def call(self, x, axis=-1, d=1.0, n=1):
        """
        Derivative in Fourier space

        Parameters
        ----------
        x : tf.Tensor
            Input array, can be either real or complex valued.
        axis : int (optional)
            Axis along which `x` is derived.
            Defaults to `-1`.
        d : float (optional)
            Discretization of `x` along the dimension of `axis`.
            Defaults to `1.0`.
        n : int (optional)
            Order of differentiation.
            Defaults to `1`.

        """

        return _fft_derivative(x=x, axis=axis, d=d, n=n)
    
def fft_derivative(x, axis=-1, d=1.0, n=1):
    """
    Derivative in Fourier space

    Parameters
    ----------
    x : tf.Tensor
        Input array, can be either real or complex valued.
    axis : int (optional)
        Axis along which `x` is derived.
        Defaults to `-1`.
    d : float (optional)
        Discretization of `x` along the dimension of `axis`.
        Defaults to `1.0`.
    n : int (optional)
        Order of differentiation.
        Defaults to `1`.

    """
    if any_symbolic_tensors(cast_to_complex(x)):
        return FFTDerivative.symbolic_call(x)
    return fft_derivative(x)
