"""Wavelet operator."""

from collections.abc import Sequence
from typing import Literal

import torch
from ptwt.conv_transform import wavedec, waverec
from ptwt.conv_transform_2 import wavedec2, waverec2
from ptwt.conv_transform_3 import wavedec3, waverec3
from pywt import Wavelet
from pywt._multilevel import _check_level

from mrpro.operators.LinearOperator import LinearOperator

# Switch off formatter to avoid having only one wavelet type per line
# fmt: off
WaveletType = Literal[
    'haar',
    'db1', 'db2', 'db3', 'db4', 'db5', 'db6', 'db7', 'db8', 'db9',
    'db10', 'db11', 'db12', 'db13', 'db14', 'db15', 'db16', 'db17', 'db18', 'db19',
    'db20', 'db21', 'db22', 'db23', 'db24', 'db25', 'db26', 'db27', 'db28', 'db29',
    'db30', 'db31', 'db32', 'db33', 'db34', 'db35', 'db36', 'db37', 'db38',
    'sym2', 'sym3', 'sym4', 'sym5', 'sym6', 'sym7', 'sym8', 'sym9',
    'sym10', 'sym11', 'sym12', 'sym13', 'sym14', 'sym15', 'sym16', 'sym17', 'sym18', 'sym19', 'sym20',
    'coif1', 'coif2', 'coif3', 'coif4', 'coif5', 'coif6', 'coif7', 'coif8', 'coif9',
    'coif10', 'coif11', 'coif12', 'coif13', 'coif14', 'coif15', 'coif16', 'coif17',
    'bior1.1', 'bior1.3', 'bior1.5', 'bior2.2', 'bior2.4', 'bior2.6', 'bior2.8',
    'bior3.1', 'bior3.3', 'bior3.5', 'bior3.7', 'bior3.9', 'bior4.4', 'bior5.5', 'bior6.8',
    'rbio1.1', 'rbio1.3', 'rbio1.5', 'rbio2.2', 'rbio2.4', 'rbio2.6', 'rbio2.8',
    'rbio3.1', 'rbio3.3', 'rbio3.5', 'rbio3.7', 'rbio3.9', 'rbio4.4', 'rbio5.5', 'rbio6.8',
    'dmey',
    'gaus1', 'gaus2', 'gaus3', 'gaus4', 'gaus5', 'gaus6', 'gaus7', 'gaus8', # noqa: typos
    'mexh',
    'morl',
    'cgau1', 'cgau2', 'cgau3', 'cgau4', 'cgau5', 'cgau6', 'cgau7', 'cgau8',
    'shan',
    'fbsp',
    'cmor'
]
class WaveletOp(LinearOperator):
    """Wavelet operator class."""

    def __init__(
        self,
        domain_shape: Sequence[int] | None = None,
        dim: tuple[int] | tuple[int, int] | tuple[int, int, int] = (-2, -1),
        wavelet_name: WaveletType = 'db4',
        level: int | None = None,
    ):
        """Wavelet operator.

        For complex images the wavelet coefficients are calculated for real and imaginary part separately.

        For a 2D image, the coefficients are labeled `[aa, (ad_n, da_n, dd_n), ..., (ad_1, da_1, dd_1)]` where a refers
        to the approximation coefficients and d to the detail coefficients. The index indicates the level.

        Parameters
        ----------
        domain_shape
            Shape of domain where wavelets are calculated. If set to `None` the shape is taken from the input of the
            forward operator. The adjoint operator will raise an error.
        dim
            Dimensions (axes) where wavelets are calculated
        wavelet_name
            Name of wavelets
        level
            Highest wavelet level. If set to `None`, the highest possible level is calculated based on the domain shape.

        Raises
        ------
        `ValueError`
            If wavelets are calculated for more than three dimensions.
        `ValueError`
            If wavelet dimensions and domain shape do not match.
        `NotImplementedError`
            If any dimension of the domain shape is odd. Adjoint will lead to the wrong domain shape.
        """
        super().__init__()
        self._domain_shape = domain_shape
        self._wavelet_name = wavelet_name
        self._dim = dim
        self._level = level

        # number of wavelet directions
        if len(dim) == 1:
            self.n_wavelet_directions = 1
        elif len(dim) == 2:
            self.n_wavelet_directions = 3
        elif len(dim) == 3:
            self.n_wavelet_directions = 7
        else:
            raise ValueError('Only 1D, 2D and 3D wavelet transforms are supported.')

        if domain_shape is not None:
            if len(dim) != len(domain_shape):
                raise ValueError(
                    'Number of dimensions along which the wavelet transform should be calculated needs to',
                    'be the same as the domain shape',
                )

            if any(d % 2 for d in domain_shape):
                raise NotImplementedError(
                    'ptwt only supports wavelet transforms for tensors with even number of '
                    'entries for all considered dimensions.'
                )

            # size of wavelets
            wavelet_length = torch.as_tensor((Wavelet(wavelet_name).dec_len,) * len(domain_shape))

            # calculate shape of wavelet coefficients at each level
            current_shape = torch.as_tensor(domain_shape)

            # if level is None, select the highest possible level.
            # raise error/warnings if level is not possible or lead to boundary effects
            verified_level = _check_level(domain_shape, wavelet_length, level)

            if verified_level == 0:  # only a/aa/aaa component possible
                self.coefficients_shape:list[torch.Size] = [torch.Size(domain_shape)]
            else:
                self.coefficients_shape = []
                for _ in range(verified_level):
                    # Add padding
                    current_shape = (current_shape / 2).ceil() + wavelet_length // 2 - 1
                    self.coefficients_shape.extend(
                        [torch.Size(current_shape.to(dtype=torch.int64))] * self.n_wavelet_directions
                    )

                self.coefficients_shape = self.coefficients_shape[::-1]
                self.coefficients_shape.insert(0, self.coefficients_shape[0])  # shape of a/aa/aaa term

    def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor,]:
        """Apply wavelet decomposition to the input tensor.

        This method calculates the wavelet coefficients of the input (image) data
        along the specified dimensions. For complex inputs, the wavelet transform
        is applied to the real and imaginary parts separately. The resulting
        coefficients from different levels and orientations are stacked along
        a single dimension.

        Parameters
        ----------
        x
            Input (image) data tensor.

        Returns
        -------
            The wavelet coefficients, stacked along one dimension.
            The first wavelet dimension specified in `dim` during initialization is used for stacking.

        Raises
        ------
        ValueError
            If the dimensions along which wavelets are to be calculated (specified
            during initialization) are not unique after normalization.
        """
        return super().__call__(x)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor,]:
        """Apply forward of WaveletOp.

        .. note::
            Prefer calling the instance of the WaveletOp operator as ``operator(x)`` over
            directly calling this method. See this PyTorch `discussion <https://discuss.pytorch.org/t/is-model-forward-x-the-same-as-model-call-x/33460/3>`_.
        """
        # normalize axes to allow negative indexing in input
        dim = tuple(d % x.ndim for d in self._dim)
        if len(dim) != len(set(dim)):
            raise ValueError(f'Axis must be unique. Normalized axis are {dim}')

        # move axes where wavelets are calculated to the end
        x = x.moveaxis(dim, list(range(-len(self._dim), 0)))

        # the ptwt functions work only for real data, thus we handle complex inputs as an additional channel
        x_real = torch.view_as_real(x).moveaxis(-1, 0) if x.is_complex() else x

        if len(self._dim) == 1:
            coeffs_1d = wavedec(x_real, self._wavelet_name, level=self._level, mode='zero', axis=-1)
            coefficients_list = self._format_coeffs_1d(coeffs_1d)
        elif len(self._dim) == 2:
            coeffs_2d = wavedec2(x_real, self._wavelet_name, level=self._level, mode='zero', axes=(-2, -1))
            coefficients_list = self._format_coeffs_2d(coeffs_2d)
        elif len(self._dim) == 3:
            coeffs_3d = wavedec3(x_real, self._wavelet_name, level=self._level, mode='zero', axes=(-3, -2, -1))
            coefficients_list = self._format_coeffs_3d(coeffs_3d)
        else:
            raise ValueError(f'Wavelets are only available for 1D, 2D and 3D and not {self._dim}D')

        # stack multi-resolution wavelets along single dimension
        coefficients_stack = self._coeff_to_stacked_tensor(coefficients_list)
        if x.is_complex():
            coefficients_stack = torch.moveaxis(
                coefficients_stack, -1, min(dim) + 1
            )  # +1 because first dim is real/imag
            coefficients_stack = torch.view_as_complex(coefficients_stack.moveaxis(0, -1).contiguous())
        else:
            coefficients_stack = torch.moveaxis(coefficients_stack, -1, min(dim))

        # move stacked coefficients to first wavelet dimension
        return (coefficients_stack,)

    def adjoint(self, coefficients_stack: torch.Tensor) -> tuple[torch.Tensor]:
        """Apply inverse wavelet transform (reconstruction) to wavelet coefficients.

        This method reconstructs the (image) data from its wavelet coefficients.
        The input `coefficients_stack` is first unstacked into a list of
        coefficients for different levels and orientations, and then the
        inverse wavelet transform is applied. For complex inputs, the transform
        is applied to real and imaginary parts separately.

        Parameters
        ----------
        coefficients_stack
            Tensor of wavelet coefficients, stacked along one dimension
            (assumed to be the first wavelet dimension specified in `dim`
            during initialization).

        Returns
        -------
            The reconstructed (image) data tensor.

        Raises
        ------
        ValueError
            If `domain_shape` was not defined during operator initialization.
        ValueError
            If the dimensions along which wavelets were calculated (specified
            during initialization) are not unique after normalization.
        """
        if self._domain_shape is None:
            raise ValueError('Adjoint requires to define the domain_shape in init()')

        # normalize axes to allow negative indexing in input
        dim = tuple(d % (coefficients_stack.ndim + len(self._dim) - 1) for d in self._dim)
        if len(dim) != len(set(dim)):
            raise ValueError(f'Axis must be unique. Normalized axis are {dim}')

        coefficients_stack = torch.moveaxis(coefficients_stack, min(dim), -1)

        # the ptwt functions work only for real data, thus we handle complex inputs as an additional channel
        coefficients_stack_real = (
            torch.view_as_real(coefficients_stack).moveaxis(-1, 0)
            if coefficients_stack.is_complex()
            else coefficients_stack
        )

        coefficients_list = self._stacked_tensor_to_coeff(coefficients_stack_real)

        if len(self._dim) == 1:
            coeffs_1d = self._undo_format_coeffs_1d(coefficients_list)
            data = waverec(coeffs_1d, self._wavelet_name, axis=-1)
        elif len(self._dim) == 2:
            coeffs_2d = self._undo_format_coeffs_2d(coefficients_list)
            data = waverec2(coeffs_2d, self._wavelet_name, axes=(-2, -1))
        elif len(self._dim) == 3:
            coeffs_3d = self._undo_format_coeffs_3d(coefficients_list)
            data = waverec3(coeffs_3d, self._wavelet_name, axes=(-3, -2, -1))
        else:
            raise ValueError(f'Wavelets are only available for 1D, 2D and 3D and not {self._dim}D')

        # undo moving of axes
        if coefficients_stack.is_complex():
            data = torch.moveaxis(
                data, list(range(-len(self._dim), 0)), [d + 1 for d in dim]
            )  # +1 because first dim is real/imag
            # if we deal with complex coefficients, we also return complex data
            data = torch.view_as_complex(data.moveaxis(0, -1).contiguous())
        else:
            data = torch.moveaxis(data, list(range(-len(self._dim), 0)), dim)

        return (data,)

    @staticmethod
    def _format_coeffs_1d(coefficients: list[torch.Tensor]) -> list[torch.Tensor]:
        """Format 1D wavelet coefficients to MRpro format.

        At the moment, this function just returns the input coefficients as is:
        [a, d_n, ..., d_1]
        """
        return coefficients

    @staticmethod
    def _undo_format_coeffs_1d(coefficients: list[torch.Tensor]) -> list[torch.Tensor]:
        """Undo format 1D wavelet coefficients to MRpro format.

        At the moment, this function just returns the input coefficients as is:
        [a, d_n, ..., d_1]
        """
        return coefficients

    @staticmethod
    def _format_coeffs_2d(
        coefficients: list[torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
    ) -> list[torch.Tensor]:
        """Format 2D wavelet coefficients to MRpro format.

        Converts from   [aa, (ad_n, da_n, dd_n), ..., (ad_1, da_1, dd_1)]
        to              [aa, ad_n, da_n, dd_n, ..., ad_1, da_1, dd_1]
        """
        coeffs_mrpro_format: list = [coefficients[0]]
        for c_tuple in coefficients[1:]:
            coeffs_mrpro_format.extend(c_tuple)
        return coeffs_mrpro_format

    def _undo_format_coeffs_2d(
        self,
        coefficients: list[torch.Tensor],
    ) -> list[torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
        """Undo format 2D wavelet coefficients to MRpro format.

        Converts from   [aa, ad_n, da_n, dd_n, ..., ad_1, da_1, dd_1]
        to              [aa, (ad_n, da_n, dd_n), ..., (ad_1, da_1, dd_1)]
        """
        coeffs_ptwt_format: list = [coefficients[0]]
        for i in range(1, len(coefficients), self.n_wavelet_directions):
            coeffs_ptwt_format.append(tuple(coefficients[i : i + self.n_wavelet_directions]))
        return coeffs_ptwt_format

    @staticmethod
    def _format_coeffs_3d(coefficients: list[torch.Tensor | dict[str, torch.Tensor]]) -> list[torch.Tensor]:
        """Format 3D wavelet coefficients to MRpro format.

        Converts from   [aaa, {aad_n, ada_n, add_n, ...}, ..., {aad_1, ada_1, add_1, ...}]
        to              [aaa, aad_n, ada_n, add_n, ..., ..., aad_1, ada_1, add_1, ...]
        """
        coeffs_mrpro_format: list = [coefficients[0]]
        for c_dict in coefficients[1:]:
            coeffs_mrpro_format.extend(c_dict.values())
        return coeffs_mrpro_format

    def _undo_format_coeffs_3d(
        self,
        coefficients: list[torch.Tensor],
    ) -> list[torch.Tensor | dict[str, torch.Tensor]]:
        """Undo format 3D wavelet coefficients to MRpro format.

        Converts from   [aaa, aad_n, ada_n, add_n, ..., ..., aad_1, ada_1, add_1, ...]
        to              [aaa, {aad_n, ada_n, add_n, ...}, ..., {aad_1, ada_1, add_1, ...}]
        """
        coeffs_ptwt_format: list = [coefficients[0]]
        for i in range(1, len(coefficients), self.n_wavelet_directions):
            coeffs_ptwt_format.append(
                dict(
                    zip(
                        ['aad', 'ada', 'add', 'daa', 'dad', 'dda', 'ddd'], # noqa: typos
                        coefficients[i : i + self.n_wavelet_directions],
                        strict=True,
                    )
                )
            )
        return coeffs_ptwt_format

    def _coeff_to_stacked_tensor(self, coefficients: list[torch.Tensor]) -> torch.Tensor:
        """Stack wavelet coefficients into 1D tensor.

        During the calculation of the of the wavelet coefficient ptwt uses padding. To ensure the wavelet operator is
        an isometry, cropping is needed.

        Parameters
        ----------
        coefficients
            wavelet coefficients in the format
            1D: [a, d_n, ..., d_1]
            2D: [aa, ad_n, da_n, dd_n, ..., ad_1, da_1, dd_1]
            3D: [aaa, aad_n, ada_n, add_n, ..., ..., aad_1, ada_1, add_1, ...]

        Returns
        -------
            stacked coefficients in tensor [...,coeff]
        """
        return torch.cat(
            [coeff.flatten(start_dim=-len(self._dim)) for coeff in coefficients]
            if len(self._dim) > 1
            else coefficients,
            dim=-1,
        )

    def _stacked_tensor_to_coeff(self, coefficients_stack: torch.Tensor) -> list[torch.Tensor]:
        """Unstack wavelet coefficients.

        Parameters
        ----------
        coefficients_stack
            stacked coefficients in tensor [...,coeff]


        Returns
        -------
            wavelet coefficients in the format
            1D: [a, d_n, ..., d_1]
            2D: [aa, ad_n, da_n, dd_n, ..., ad_1, da_1, dd_1]
            3D: [aaa, aad_n, ada_n, add_n, ..., ..., aad_1, ada_1, add_1, ...]
        """
        coefficients = coefficients_stack.split(
            [shape.numel() for shape in self.coefficients_shape], dim=-1
        )
        return [
            coeff.reshape((*coeff.shape[:-1], *shape))
            for coeff, shape in zip(coefficients, self.coefficients_shape, strict=True)
        ]
