# (C) Copyright 2025 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.


import abc
import logging

import einops
import torch
import torch.fft

from anemoi.training.utils.enums import TensorDim

LOGGER = logging.getLogger(__name__)


class SpectralTransform(torch.nn.Module):
    """Abstract base class for spectral transforms."""

    @abc.abstractmethod
    def forward(
        self,
        data: torch.Tensor,
    ) -> torch.Tensor:
        """Transform data to spectral domain.

        Parameters
        ----------
        data : torch.Tensor
            Input data in the spatial domain of expected shape
            `[batch, ensemble, points, variables]`.

        Returns
        -------
        torch.Tensor
            Data transformed to the spectral domain, of shape
            `[batch, ensemble, y_freq, x_freq, variables]`.
        """


class FFT2D(SpectralTransform):
    """2D Fast Fourier Transform (FFT) implementation."""

    def __init__(
        self,
        x_dim: int,
        y_dim: int,
        nodes_slice: tuple[int, int] = (0, None),
    ) -> None:
        """2D FFT Transform.

        Parameters
        ----------
        x_dim : int
            size of the spatial dimension x of the original data in 2D
        y_dim : int
            size of the spatial dimension y of the original data in 2D
        """
        self.x_dim = x_dim
        self.y_dim = y_dim
        self.nodes_slice = slice(*nodes_slice)

    def forward(
        self,
        data: torch.Tensor,
    ) -> torch.Tensor:
        data = torch.index_select(
            data, TensorDim.GRID, torch.arange(*self.nodes_slice.indices(data.size(TensorDim.GRID)))
        )
        var = data.shape[-1]
        try:
            data = einops.rearrange(data, "... (y x) v -> ... y x v", x=self.x_dim, y=self.y_dim, v=var)
        except Exception as e:
            raise einops.EinopsError(
                f"Possible dimension mismatch in einops.rearrange in FFT2D layer: "
                f"expected (y * x) == last spatial dim with y={self.y_dim}, x={self.x_dim}"
            ) from e
        return torch.fft.fft2(data, dim=(-2, -3))


class SHT(SpectralTransform):
    """Placeholder for Spherical Harmonics Transform."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__()

    def forward(
        self,
        data: torch.Tensor,
    ) -> torch.Tensor:
        """Transform data to spectral domain using spherical harmonics.

        Parameters
        ----------
        data : torch.Tensor
            Input data in the spatial domain.

        Returns
        -------
        torch.Tensor
            Data transformed to the spectral domain.
        """
        msg = "Spherical harmonics transform is not implemented yet."
        raise NotImplementedError(msg)
