r"""PyWavelets-compatible Discrete Wavelet Transform implementations.

This module provides DWT implementations that exactly match PyWavelets behavior
while maintaining full gradient support for PyTorch.

Classes
-------
DWT1D
    1D Discrete Wavelet Transform with multi-level support.
DWT2D
    2D Discrete Wavelet Transform using separable 1D transforms.

Functions
---------
get_wavelet_filters(wavelet_name)
    Extract filter coefficients from PyWavelets.

Examples
--------
Basic 1D wavelet transform:

>>> import torch
>>> from spectrans.transforms.wavelet import DWT1D
>>> dwt = DWT1D(wavelet='db4', levels=2)
>>> x = torch.randn(32, 256)
>>> cA, cD_list = dwt.decompose(x)
>>> x_rec = dwt.reconstruct((cA, cD_list))
>>> error = torch.max(torch.abs(x - x_rec))
>>> print(f"Reconstruction error: {error:.2e}")  # Should be < 1e-6

2D wavelet transform for images:

>>> from spectrans.transforms.wavelet import DWT2D
>>> dwt2d = DWT2D(wavelet='db2', levels=2)
>>> image = torch.randn(1, 64, 64)
>>> ll, detail_bands = dwt2d.decompose(image)
>>> reconstructed = dwt2d.reconstruct((ll, detail_bands))

Multi-level decomposition with energy analysis:

>>> dwt = DWT1D(wavelet='db4', levels=3)
>>> x = torch.randn(1, 512)
>>> cA, cD_list = dwt.decompose(x)
>>> # Verify Parseval's theorem for orthogonal wavelets
>>> energy_input = torch.sum(x ** 2)
>>> energy_coeffs = torch.sum(cA ** 2) + sum(torch.sum(cD ** 2) for cD in cD_list)
>>> print(f"Energy ratio: {energy_coeffs / energy_input:.6f}")  # Should be ≈ 1.0

Notes
-----
Mathematical Foundations:
The Discrete Wavelet Transform (DWT) decomposes a signal into approximation
and detail coefficients through iterative filtering and downsampling.

For a signal $\mathbf{x}[n]$ of length $N$, the single-level DWT produces:

$$
c_A[k] = \sum_{n} h[n-2k] \cdot \mathbf{x}[n]
$$

$$
c_D[k] = \sum_{n} g[n-2k] \cdot \mathbf{x}[n]
$$

Where $h[n]$ and $g[n]$ are the low-pass and high-pass analysis
filters. The reconstruction is achieved through:

$$
\mathbf{x}[n] = \sum_{k} h'[n-2k] \cdot c_A[k] + \sum_{k} g'[n-2k] \cdot c_D[k]
$$

Multi-Resolution Analysis:
The $J$-level DWT recursively applies the transform to approximation
coefficients, creating a dyadic decomposition where each level $j$ has
length $N/2^j$ and frequency band $[0, \pi/2^j]$ for approximations.

Perfect Reconstruction:
For orthogonal wavelets: $h'[n] = h[-n]$ and $g'[n] = g[-n]$.
The transform preserves energy: $\|\mathbf{x}\|^2 = \|\mathbf{c}_A\|^2 + \sum_{j} \|\mathbf{c}_{D_j}\|^2$

Implementation Details:

- Convolution starts at index $(\text{step} - 1) = 1$ for stride 2
- Symmetric mode reflects without edge repeat: ``[d,c,b,a | a,b,c,d | d,c,b,a]``
- Uses ``conv1d`` with flipped filters for correlation
- IDWT uses ``conv_transpose1d`` with stride 2 for implicit upsampling
- Output lengths follow PyWavelets formulas

Algorithm Complexity:

- Forward/Inverse DWT: $O(N)$ for $N$-length signal
- Memory: $O(N)$ for coefficients

Gradient Support:
All operations use native PyTorch operations ensuring full autograd support.

Numerical Precision:

- Filters use ``float64`` for extraction, ``float32`` for computation
- Perfect reconstruction to $\sim 10^{-7}$ for ``float32``

Supported Wavelets:
Daubechies (``db1``-``db38``), Symlets (``sym2``-``sym20``),
Coiflets (``coif1``-``coif17``), Biorthogonal (``bior``/``rbio``),
Discrete Meyer (``dmey``), Haar (``haar``)

References
----------
Stéphane Mallat. 2009. A Wavelet Tour of Signal Processing: The Sparse Way,
3rd edition. Academic Press, Boston.

Ingrid Daubechies. 1992. Ten Lectures on Wavelets. SIAM, Philadelphia.

Gilbert Strang and Truong Nguyen. 1996. Wavelets and Filter Banks.
Wellesley-Cambridge Press, Wellesley.

PyWavelets Development Team. 2024. PyWavelets: Wavelet transforms in Python.
https://pywavelets.readthedocs.io/

See Also
--------
spectrans.transforms.base : Base transform interfaces
spectrans.layers.mixing.wavelet : Wavelet mixing layers
spectrans.transforms.fourier : Fourier transform implementations
"""

import pywt
import torch
import torch.nn.functional as F

from ..core.registry import register_component
from ..core.types import Tensor, WaveletType
from ..utils.padding import wavelet_symmetric_pad
from .base import MultiResolutionTransform, MultiResolutionTransform2D


def get_wavelet_filters(wavelet_name: str) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    """Get filter coefficients from PyWavelets.

    Parameters
    ----------
    wavelet_name : str
        Name of the wavelet (e.g., 'db1', 'db2', 'db4', 'sym2').

    Returns
    -------
    tuple[Tensor, Tensor, Tensor, Tensor]
        Tuple of (dec_lo, dec_hi, rec_lo, rec_hi) filter tensors.

    Raises
    ------
    ValueError
        If wavelet is not supported by PyWavelets.
    """
    try:
        wavelet = pywt.Wavelet(wavelet_name)
    except ValueError as e:
        msg = f"Unsupported wavelet: {wavelet_name}"
        raise ValueError(msg) from e

    # Extract filters exactly as PyWavelets provides them
    # Use float64 for maximum precision compatibility with PyWavelets
    dec_lo = torch.tensor(wavelet.dec_lo, dtype=torch.float64)
    dec_hi = torch.tensor(wavelet.dec_hi, dtype=torch.float64)
    rec_lo = torch.tensor(wavelet.rec_lo, dtype=torch.float64)
    rec_hi = torch.tensor(wavelet.rec_hi, dtype=torch.float64)

    return dec_lo, dec_hi, rec_lo, rec_hi


@register_component("transform", "dwt1d")
class DWT1D(MultiResolutionTransform):
    """PyWavelets-compatible 1D Discrete Wavelet Transform.

    This implementation exactly matches PyWavelets behavior based on
    comprehensive C code analysis. It supports multi-level decomposition
    and achieves perfect reconstruction (< 1e-6 error) for all wavelets.

    Parameters
    ----------
    wavelet : WaveletType, default='db4'
        Wavelet type (e.g., 'db1', 'db2', 'db4', 'db8', 'sym2', 'coif1').
    levels : int, default=1
        Number of decomposition levels.
    mode : str, default='symmetric'
        Boundary handling mode (currently only 'symmetric' supported).

    Attributes
    ----------
    wavelet : str
        The wavelet type being used.
    levels : int
        Number of decomposition levels.
    mode : str
        Boundary handling mode.
    dec_lo : Tensor
        Low-pass decomposition filter.
    dec_hi : Tensor
        High-pass decomposition filter.
    rec_lo : Tensor
        Low-pass reconstruction filter.
    rec_hi : Tensor
        High-pass reconstruction filter.
    filter_length : int
        Length of the wavelet filters.

    Examples
    --------
    >>> dwt = DWT1D(wavelet='db4', levels=3)
    >>> x = torch.randn(16, 256)  # batch_size=16, length=256
    >>> cA, cD_list = dwt.decompose(x)
    >>> print(f"Approximation shape: {cA.shape}")
    >>> print(f"Number of detail levels: {len(cD_list)}")
    >>> x_rec = dwt.reconstruct((cA, cD_list))
    >>> error = torch.max(torch.abs(x - x_rec))
    >>> print(f"Reconstruction error: {error:.2e}")
    """

    # Explicit type annotations for buffers (mypy doesn't understand register_buffer)
    dec_lo: Tensor
    dec_hi: Tensor
    rec_lo: Tensor
    rec_hi: Tensor

    def __init__(self, wavelet: WaveletType = "db4", levels: int = 1, mode: str = "symmetric"):
        super().__init__(levels=levels)
        self.wavelet = wavelet
        self.mode = mode

        if mode != "symmetric":
            msg = f"Mode '{mode}' not yet supported. Only 'symmetric' is implemented."
            raise NotImplementedError(msg)

        # Get filter coefficients from PyWavelets
        dec_lo, dec_hi, rec_lo, rec_hi = get_wavelet_filters(wavelet)

        # Register as buffers (not parameters) since they're fixed
        # Convert to float32 for efficiency in neural networks
        self.register_buffer("dec_lo", dec_lo.float())
        self.register_buffer("dec_hi", dec_hi.float())
        self.register_buffer("rec_lo", rec_lo.float())
        self.register_buffer("rec_hi", rec_hi.float())

        self.filter_length = len(dec_lo)

    def _pad_symmetric(self, x: Tensor, pad_len: int, dim: int = -1) -> Tensor:
        """Apply symmetric padding for wavelet transforms.

        Uses the wavelet_symmetric_pad function from utils.padding which
        implements the exact PyWavelets symmetric padding mode WITH edge repeat.
        For signal [a,b,c,d] with pad=2, creates [b,a|a,b,c,d|d,c].

        Parameters
        ----------
        x : Tensor
            Input tensor to pad.
        pad_len : int
            Number of samples to pad on each side.
        dim : int
            Dimension to pad along.

        Returns
        -------
        Tensor
            Padded tensor.
        """
        if pad_len <= 0:
            return x
        return wavelet_symmetric_pad(x, pad_len, dim)

    def _single_dwt(self, x: Tensor, dim: int = -1) -> tuple[Tensor, Tensor]:
        """Single-level DWT matching PyWavelets exactly.

        Critical implementation details from C code:
        1. Start convolution at index (step-1) = 1
        2. Use symmetric padding (reflection without edge)
        3. Apply filters as provided by PyWavelets

        Parameters
        ----------
        x : Tensor
            Input signal.
        dim : int
            Dimension to transform along.

        Returns
        -------
        tuple[Tensor, Tensor]
            Tuple of (cA, cD) coefficients.
        """
        # Handle different tensor dimensions
        if x.ndim == 1:
            # Add batch dimension
            x = x.unsqueeze(0)
            cA, cD = self._single_dwt(x, dim=-1)
            return cA.squeeze(0), cD.squeeze(0)

        # Apply symmetric padding
        pad_len = self.filter_length - 1
        x_padded = self._pad_symmetric(x, pad_len, dim=dim)

        # Prepare for convolution
        if dim == -1 or dim == x.ndim - 1:
            # Last dimension case (most common)
            x_padded = x_padded.unsqueeze(1)  # Add channel dimension

            # CRITICAL: Start from position (step - 1) = 1
            # This is the key insight from PyWavelets C code
            x_conv = x_padded[:, :, 1:] if pad_len > 0 else x_padded

            # Apply filters (flip for correlation -> convolution)
            # Ensure filters match input tensor's dtype
            h_filter = self.dec_lo.to(x_conv.dtype).flip(0).unsqueeze(0).unsqueeze(0)
            g_filter = self.dec_hi.to(x_conv.dtype).flip(0).unsqueeze(0).unsqueeze(0)

            # Convolve with stride 2 (downsampling)
            cA = F.conv1d(x_conv, h_filter, stride=2).squeeze(1)
            cD = F.conv1d(x_conv, g_filter, stride=2).squeeze(1)

            # Ensure correct output length (PyWavelets formula)
            signal_len = x.shape[-1]
            expected_len = (signal_len + self.filter_length - 1) // 2
            cA = cA[..., :expected_len]
            cD = cD[..., :expected_len]

        else:
            # Handle arbitrary dimension
            # Move the target dimension to last position
            x_transposed = x.transpose(dim, -1)
            cA, cD = self._single_dwt(x_transposed, dim=-1)
            # Move dimension back
            cA = cA.transpose(dim, -1)
            cD = cD.transpose(dim, -1)

        return cA, cD

    def _single_dwt_nd(self, x: Tensor, axis: int) -> tuple[Tensor, Tensor]:
        """Apply 1D DWT along any axis of n-dimensional tensor.

        This method handles n-dimensional tensors by reshaping them to 2D,
        applying the DWT, and reshaping back. This follows PyWavelets'
        approach of having a dwt_axis function for n-dimensional arrays.

        Parameters
        ----------
        x : Tensor
            Input tensor of any dimensionality.
        axis : int
            Axis along which to apply the DWT.

        Returns
        -------
        tuple[Tensor, Tensor]
            Tuple of (cA, cD) coefficients with same dimensionality as input
            except along the transform axis which is reduced by downsampling.
        """
        # Handle negative axis
        if axis < 0:
            axis = x.ndim + axis

        # Move target axis to last position
        x_moved = x.moveaxis(axis, -1)

        # Store original shape
        original_shape = x_moved.shape
        batch_shape = original_shape[:-1]
        signal_len = original_shape[-1]

        # Reshape to 2D (batch, signal_len)
        x_2d = x_moved.reshape(-1, signal_len)

        # Apply standard 1D DWT (which expects 2D tensors)
        cA, cD = self._single_dwt(x_2d, dim=-1)

        # Calculate new shape after DWT
        new_signal_len = cA.shape[-1]
        new_shape = (*batch_shape, new_signal_len)

        # Reshape back to original dimensionality
        cA = cA.reshape(new_shape).moveaxis(-1, axis)
        cD = cD.reshape(new_shape).moveaxis(-1, axis)

        return cA, cD

    def _single_idwt(self, cA: Tensor, cD: Tensor, dim: int = -1) -> Tensor:
        """Single-level inverse DWT in pure PyTorch using transpose convolution.

        Based on the PyWavelets algorithm, uses transpose convolution with stride 2
        for implicit upsampling, matching the approach from pywavelets-implementation-plan.md.

        Parameters
        ----------
        cA : Tensor
            Approximation coefficients.
        cD : Tensor
            Detail coefficients.
        dim : int
            Dimension to reconstruct along.

        Returns
        -------
        Tensor
            Reconstructed signal.
        """
        # Ensure tensors have same shape
        if cA.shape != cD.shape:
            raise ValueError(
                f"Coefficients must have same shape. Got cA: {cA.shape}, cD: {cD.shape}"
            )

        # Move dimension to last for processing
        if dim != -1 and dim != cA.ndim - 1:
            cA = cA.transpose(dim, -1)
            cD = cD.transpose(dim, -1)

        batch_shape = cA.shape[:-1]
        coeffs_len = cA.shape[-1]

        # Calculate expected output length using PyWavelets formula
        # For symmetric mode: output_len = 2 * coeffs_len - filter_len + 2
        output_len = 2 * coeffs_len - self.filter_length + 2

        # Reshape for conv_transpose1d: [batch, 1, coeffs_len]
        cA_reshaped = cA.reshape(-1, 1, coeffs_len)
        cD_reshaped = cD.reshape(-1, 1, coeffs_len)

        # Prepare filters for transpose convolution
        # Filters should NOT be flipped for transpose convolution
        # Ensure filters match input tensor's dtype
        rec_lo_filter = self.rec_lo.to(cA_reshaped.dtype).unsqueeze(0).unsqueeze(0)
        rec_hi_filter = self.rec_hi.to(cD_reshaped.dtype).unsqueeze(0).unsqueeze(0)

        # Apply transpose convolution with stride 2 (implicit upsampling)
        # This effectively inserts zeros between coefficients and convolves
        rec_from_cA = F.conv_transpose1d(cA_reshaped, rec_lo_filter, stride=2)
        rec_from_cD = F.conv_transpose1d(cD_reshaped, rec_hi_filter, stride=2)

        # Align outputs based on filter length
        # The output from transpose convolution needs to be trimmed
        # to match PyWavelets alignment
        # For Haar (filter_len=2): no trimming from start
        # For longer filters: trim (filter_len - 2) from start
        if self.filter_length == 2:
            # Haar wavelet - no trimming needed from start
            pass
        else:
            # For longer filters, trim from the beginning
            # This aligns the reconstruction with PyWavelets
            trim_start = self.filter_length - 2
            rec_from_cA = rec_from_cA[:, :, trim_start:]
            rec_from_cD = rec_from_cD[:, :, trim_start:]

        # Sum the reconstructions (as in PyWavelets)
        reconstructed = rec_from_cA + rec_from_cD

        # Reshape back to original batch shape
        reconstructed = reconstructed.reshape(*batch_shape, -1)

        # Trim to expected output length
        if reconstructed.shape[-1] != output_len:
            actual_len = reconstructed.shape[-1]
            if actual_len > output_len:
                # Trim excess samples from the end
                reconstructed = reconstructed[..., :output_len]
            else:
                raise RuntimeError(
                    f"Reconstruction length mismatch. Expected {output_len}, "
                    f"got {actual_len}. Filter length: {self.filter_length}"
                )

        # Move dimension back if it was transposed
        if dim != -1 and dim != cA.ndim - 1:
            reconstructed = reconstructed.transpose(-1, dim)

        return reconstructed

    def _single_idwt_nd(self, cA: Tensor, cD: Tensor, axis: int) -> Tensor:
        """Apply 1D inverse DWT along any axis of n-dimensional tensor.

        This method handles n-dimensional coefficient tensors by reshaping
        them to 2D, applying the inverse DWT, and reshaping back.

        Parameters
        ----------
        cA : Tensor
            Approximation coefficients of any dimensionality.
        cD : Tensor
            Detail coefficients of any dimensionality.
        axis : int
            Axis along which to apply the inverse DWT.

        Returns
        -------
        Tensor
            Reconstructed tensor with same dimensionality as input
            except along the transform axis which is upsampled.
        """
        # Handle negative axis
        if axis < 0:
            axis = cA.ndim + axis

        # Move target axis to last position
        cA_moved = cA.moveaxis(axis, -1)
        cD_moved = cD.moveaxis(axis, -1)

        # Store original shape
        original_shape = cA_moved.shape
        batch_shape = original_shape[:-1]
        coeffs_len = original_shape[-1]

        # Reshape to 2D (batch, coeffs_len)
        cA_2d = cA_moved.reshape(-1, coeffs_len)
        cD_2d = cD_moved.reshape(-1, coeffs_len)

        # Apply standard 1D inverse DWT
        reconstructed = self._single_idwt(cA_2d, cD_2d, dim=-1)

        # Calculate new shape after reconstruction
        new_signal_len = reconstructed.shape[-1]
        new_shape = (*batch_shape, new_signal_len)

        # Reshape back to original dimensionality
        reconstructed = reconstructed.reshape(new_shape).moveaxis(-1, axis)

        return reconstructed

    def _single_idwt_pytorch(self, cA: Tensor, cD: Tensor, dim: int = -1) -> Tensor:
        """Single-level inverse DWT matching PyWavelets.

        Uses transpose convolution for upsampling and reconstruction.

        Parameters
        ----------
        cA : Tensor
            Approximation coefficients.
        cD : Tensor
            Detail coefficients.
        dim : int
            Dimension to reconstruct along.

        Returns
        -------
        Tensor
            Reconstructed signal.
        """
        # Handle 1D case
        if cA.ndim == 1:
            cA = cA.unsqueeze(0)
            cD = cD.unsqueeze(0)
            result = self._single_idwt(cA, cD, dim=-1)
            return result.squeeze(0)

        if dim == -1 or dim == cA.ndim - 1:
            # Last dimension case
            cA_reshaped = cA.unsqueeze(1)
            cD_reshaped = cD.unsqueeze(1)

            # Reconstruction filters (no flip needed for transpose conv)
            # Ensure filters match input tensor's dtype
            rec_lo_filter = self.rec_lo.to(cA_reshaped.dtype).unsqueeze(0).unsqueeze(0)
            rec_hi_filter = self.rec_hi.to(cD_reshaped.dtype).unsqueeze(0).unsqueeze(0)

            # Apply transpose convolution (upsampling by stride 2)
            rec_from_cA = F.conv_transpose1d(cA_reshaped, rec_lo_filter, stride=2)
            rec_from_cD = F.conv_transpose1d(cD_reshaped, rec_hi_filter, stride=2)

            # Sum the reconstructions (as PyWavelets does)
            reconstructed = (rec_from_cA + rec_from_cD).squeeze(1)

            # Calculate expected output length
            # PyWavelets formula: 2*coeffs_len - filter_len + 2
            coeffs_len = cA.shape[-1]
            expected_len = 2 * coeffs_len - self.filter_length + 2

            # Trim to expected length
            reconstructed = reconstructed[..., :expected_len]

        else:
            # Handle arbitrary dimension
            cA_transposed = cA.transpose(dim, -1)
            cD_transposed = cD.transpose(dim, -1)
            reconstructed = self._single_idwt(cA_transposed, cD_transposed, dim=-1)
            reconstructed = reconstructed.transpose(dim, -1)

        return reconstructed

    def decompose(
        self, x: Tensor, levels: int | None = None, dim: int = -1
    ) -> tuple[Tensor, list[Tensor]]:
        """Multi-level DWT decomposition.

        Recursively applies DWT to approximation coefficients.

        Parameters
        ----------
        x : Tensor
            Input signal.
        levels : int | None
            Number of levels. If None, uses self.levels.
        dim : int
            Dimension to decompose along.

        Returns
        -------
        tuple[Tensor, list[Tensor]]
            Tuple of (approximation, [detail_1, ..., detail_N])
            where details are ordered from finest to coarsest.
        """
        if levels is None:
            levels = self.levels

        current = x
        details = []

        # Apply DWT recursively
        for _ in range(levels):
            cA, cD = self._single_dwt(current, dim=dim)
            details.append(cD)
            current = cA

        return current, details

    def reconstruct(
        self, coeffs: tuple[Tensor, list[Tensor]], dim: int = -1, output_len: int | None = None
    ) -> Tensor:
        """Multi-level DWT reconstruction.

        Parameters
        ----------
        coeffs : tuple[Tensor, list[Tensor]]
            Tuple of (approximation, [detail_1, ..., detail_N]).
        dim : int
            Dimension to reconstruct along.
        output_len : int | None
            Desired output length. If provided, the reconstructed signal
            will be trimmed or padded to this length.

        Returns
        -------
        Tensor
            Reconstructed signal.
        """
        cA, details = coeffs
        current = cA

        # Reconstruct from coarsest to finest (reverse order)
        for _i, cD in enumerate(reversed(details)):
            # For multi-level, we need to handle size mismatches
            # The reconstructed signal from a coarser level may be slightly
            # longer than the detail coefficients from the finer level

            if current.shape[dim] > cD.shape[dim]:
                # Trim current to match cD size
                # This happens because IDWT can produce slightly longer output
                target_len = cD.shape[dim]
                if dim == -1 or dim == current.ndim - 1:
                    current = current[..., :target_len]
                elif dim == 0:
                    current = current[:target_len]
                else:
                    # General case
                    slices = [slice(None)] * current.ndim
                    slices[dim] = slice(0, target_len)
                    current = current[tuple(slices)]
            elif current.shape[dim] < cD.shape[dim]:
                # This shouldn't happen with correct decomposition
                raise ValueError(
                    f"Reconstructed signal smaller than detail coefficients. "
                    f"current shape: {current.shape}, cD shape: {cD.shape}"
                )

            # Now they should have matching sizes
            current = self._single_idwt(current, cD, dim=dim)

        # Trim to desired output length if specified
        if output_len is not None and current.shape[dim] != output_len:
            if dim == -1 or dim == current.ndim - 1:
                current = current[..., :output_len]
            elif dim == 0:
                current = current[:output_len]
            else:
                slices = [slice(None)] * current.ndim
                slices[dim] = slice(0, output_len)
                current = current[tuple(slices)]

        return current


@register_component("transform", "dwt2d")
class DWT2D(MultiResolutionTransform2D):
    """PyWavelets-compatible 2D Discrete Wavelet Transform.

    Implements 2D DWT using separable 1D transforms, applying DWT
    along each dimension sequentially. Returns coefficients in the
    standard format: (LL, [(LH, HL, HH) per level]).

    Parameters
    ----------
    wavelet : WaveletType, default='db4'
        Wavelet type to use.
    levels : int, default=1
        Number of decomposition levels.
    mode : str, default='symmetric'
        Boundary handling mode.

    Attributes
    ----------
    wavelet : str
        The wavelet type.
    levels : int
        Number of decomposition levels.
    mode : str
        Boundary handling mode.
    dwt1d : DWT1D
        1D DWT instance used for separable transforms.

    Examples
    --------
    >>> dwt2d = DWT2D(wavelet='db2', levels=2)
    >>> image = torch.randn(4, 64, 64)  # batch of 4 images
    >>> ll, detail_bands = dwt2d.decompose(image)
    >>> print(f"LL shape: {ll.shape}")
    >>> for i, (lh, hl, hh) in enumerate(detail_bands):
    ...     print(f"Level {i+1} - LH: {lh.shape}, HL: {hl.shape}, HH: {hh.shape}")
    >>> reconstructed = dwt2d.reconstruct((ll, detail_bands))
    """

    def __init__(self, wavelet: WaveletType = "db4", levels: int = 1, mode: str = "symmetric"):
        super().__init__(levels=levels)
        self.wavelet = wavelet
        self.mode = mode

        # Use 1D DWT for separable 2D transform
        self.dwt1d = DWT1D(wavelet=wavelet, levels=1, mode=mode)

    def _single_level_2d(
        self, x: Tensor, dim: tuple[int, int] = (-2, -1)
    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        """Single-level 2D DWT decomposition.

        Parameters
        ----------
        x : Tensor
            2D input tensor.
        dim : tuple[int, int]
            Dimensions to decompose along.

        Returns
        -------
        tuple[Tensor, Tensor, Tensor, Tensor]
            Tuple of (LL, LH, HL, HH) coefficients following PyWavelets convention:
            - LL: approximation on both axes (aa)
            - LH: approximation on rows, detail on columns (ad)
            - HL: detail on rows, approximation on columns (da)
            - HH: detail on both axes (dd)
        """
        # Apply 1D DWT along first dimension (rows)
        # This gives us approximation and detail along the row axis
        cA_row, cD_row = self.dwt1d._single_dwt_nd(x, axis=dim[0])

        # Apply 1D DWT along second dimension (columns) to row approximation
        # ll = approx on rows, approx on cols (aa)
        # lh = approx on rows, detail on cols (ad)
        ll, lh = self.dwt1d._single_dwt_nd(cA_row, axis=dim[1])

        # Apply 1D DWT along second dimension (columns) to row detail
        # hl = detail on rows, approx on cols (da)
        # hh = detail on rows, detail on cols (dd)
        hl, hh = self.dwt1d._single_dwt_nd(cD_row, axis=dim[1])

        return ll, lh, hl, hh

    def _single_level_2d_reconstruct(
        self, ll: Tensor, lh: Tensor, hl: Tensor, hh: Tensor, dim: tuple[int, int] = (-2, -1)
    ) -> Tensor:
        """Single-level 2D DWT reconstruction.

        Parameters
        ----------
        ll, lh, hl, hh : Tensor
            2D wavelet coefficients. ll can be larger than the detail coefficients
            in multi-level reconstruction due to IDWT producing slightly longer output.
        dim : tuple[int, int]
            Dimensions to reconstruct along.

        Returns
        -------
        Tensor
            Reconstructed 2D tensor.
        """
        # Handle size mismatches in multi-level reconstruction
        # The LL (approximation) from a coarser level may be slightly larger
        # than the detail coefficients from a finer level

        # First check and trim ll to match detail coefficient sizes if needed
        if ll.shape[dim[0]] > hl.shape[dim[0]]:
            # Trim ll along first dimension to match hl
            target_size_0 = hl.shape[dim[0]]
            if dim[0] == -2 or dim[0] == ll.ndim - 2:
                ll = ll[..., :target_size_0, :]
            else:
                slices = [slice(None)] * ll.ndim
                slices[dim[0]] = slice(0, target_size_0)
                ll = ll[tuple(slices)]

        if ll.shape[dim[1]] > lh.shape[dim[1]]:
            # Trim ll along second dimension to match lh
            target_size_1 = lh.shape[dim[1]]
            if dim[1] == -1 or dim[1] == ll.ndim - 1:
                ll = ll[..., :target_size_1]
            else:
                slices = [slice(None)] * ll.ndim
                slices[dim[1]] = slice(0, target_size_1)
                ll = ll[tuple(slices)]

        # Also need to ensure hl matches hh in the second dimension
        if hl.shape[dim[1]] > hh.shape[dim[1]]:
            target_size_1 = hh.shape[dim[1]]
            if dim[1] == -1 or dim[1] == hl.ndim - 1:
                hl = hl[..., :target_size_1]
            else:
                slices = [slice(None)] * hl.ndim
                slices[dim[1]] = slice(0, target_size_1)
                hl = hl[tuple(slices)]

        # Reconstruct along second dimension first using n-dimensional method
        cA_recon = self.dwt1d._single_idwt_nd(ll, lh, axis=dim[1])
        cD_recon = self.dwt1d._single_idwt_nd(hl, hh, axis=dim[1])

        # Handle size mismatch after first reconstruction
        if cA_recon.shape[dim[0]] > cD_recon.shape[dim[0]]:
            target_size = cD_recon.shape[dim[0]]
            if dim[0] == -2 or dim[0] == cA_recon.ndim - 2:
                cA_recon = cA_recon[..., :target_size, :]
            else:
                slices = [slice(None)] * cA_recon.ndim
                slices[dim[0]] = slice(0, target_size)
                cA_recon = cA_recon[tuple(slices)]
        elif cD_recon.shape[dim[0]] > cA_recon.shape[dim[0]]:
            target_size = cA_recon.shape[dim[0]]
            if dim[0] == -2 or dim[0] == cD_recon.ndim - 2:
                cD_recon = cD_recon[..., :target_size, :]
            else:
                slices = [slice(None)] * cD_recon.ndim
                slices[dim[0]] = slice(0, target_size)
                cD_recon = cD_recon[tuple(slices)]

        # Reconstruct along first dimension using n-dimensional method
        reconstructed = self.dwt1d._single_idwt_nd(cA_recon, cD_recon, axis=dim[0])

        return reconstructed

    def decompose(
        self, x: Tensor, levels: int | None = None, dim: tuple[int, int] = (-2, -1)
    ) -> tuple[Tensor, list[tuple[Tensor, Tensor, Tensor]]]:
        """Multi-level 2D DWT decomposition.

        Parameters
        ----------
        x : Tensor
            Input 2D tensor.
        levels : int | None
            Number of levels. If None, uses self.levels.
        dim : tuple[int, int]
            Dimensions to decompose along.

        Returns
        -------
        tuple[Tensor, list[tuple[Tensor, Tensor, Tensor]]]
            Tuple of (LL, [(HL, LH, HH) per level]) following PyWavelets convention
            where HL is horizontal detail, LH is vertical detail, HH is diagonal detail.
        """
        if levels is None:
            levels = self.levels

        current = x
        detail_bands = []

        for _ in range(levels):
            ll, lh, hl, hh = self._single_level_2d(current, dim=dim)
            # PyWavelets returns (cH, cV, cD) = (HL, LH, HH)
            # So we append (HL, LH, HH) to match
            detail_bands.append((hl, lh, hh))
            current = ll

        return current, detail_bands

    def reconstruct(
        self,
        coeffs: tuple[Tensor, list[tuple[Tensor, Tensor, Tensor]]],
        dim: tuple[int, int] = (-2, -1),
    ) -> Tensor:
        """Multi-level 2D DWT reconstruction.

        Parameters
        ----------
        coeffs : tuple[Tensor, list[tuple[Tensor, Tensor, Tensor]]]
            Tuple of (LL, [(HL, LH, HH) per level]) following PyWavelets convention.
        dim : tuple[int, int]
            Dimensions to reconstruct along.

        Returns
        -------
        Tensor
            Reconstructed 2D tensor.
        """
        ll, detail_bands = coeffs
        current = ll

        # Reconstruct from coarsest to finest
        # detail_bands contains (HL, LH, HH) tuples following PyWavelets convention
        for hl, lh, hh in reversed(detail_bands):
            current = self._single_level_2d_reconstruct(current, lh, hl, hh, dim=dim)

        return current


__all__ = [
    "DWT1D",
    "DWT2D",
    "get_wavelet_filters",
]
