"""Spectral feature maps for Hilbert-space learning.

The paper emphasizes Fourier, wavelet, scattering, and learnable spectral
representations. This module implements compact CPU-friendly versions for the
8x8 Digits benchmark.
"""

from __future__ import annotations

import numpy as np
import torch


def _as_images(X: np.ndarray) -> np.ndarray:
    """Return images as (n, 8, 8) float32 in [0, 1]."""
    X = np.asarray(X, dtype=np.float32)
    if X.ndim == 2 and X.shape[1] == 64:
        X = X.reshape(-1, 8, 8)
    elif X.ndim != 3:
        raise ValueError(f"Expected (n,64) or (n,8,8), got {X.shape}")
    # sklearn digits are in [0,16]. If already normalized this is harmless.
    if X.min() >= 0.0 and X.max() > 1.0:
        X = X / 16.0
    return X.astype(np.float32)


def pixel_features(X: np.ndarray) -> np.ndarray:
    """Identity Hilbert embedding over pixels."""
    imgs = _as_images(X)
    return imgs.reshape(imgs.shape[0], -1).astype(np.float32)


def fourier_features(X: np.ndarray) -> np.ndarray:
    """Real-valued 2D Fourier feature map.

    We concatenate real and imaginary FFT coefficients. This is a spectral
    embedding in a global orthonormal basis.
    """
    imgs = _as_images(X)
    fft = np.fft.fft2(imgs, axes=(-2, -1), norm="ortho")
    feat = np.concatenate([fft.real.reshape(len(imgs), -1), fft.imag.reshape(len(imgs), -1)], axis=1)
    return feat.astype(np.float32)


def haar2d_single(img: np.ndarray, levels: int = 3) -> np.ndarray:
    """Orthonormal 2D Haar wavelet transform for square power-of-two images.

    Coefficients are packed in-place using the standard quadrant layout.
    For 8x8 images, levels=3 yields a full multiresolution transform.
    """
    x = np.asarray(img, dtype=np.float32).copy()
    h, w = x.shape
    if h != w or (h & (h - 1)) != 0:
        raise ValueError("haar2d_single expects square power-of-two images")
    size = h
    for _ in range(levels):
        if size < 2:
            break
        block = x[:size, :size]
        a = block[0::2, 0::2]
        b = block[0::2, 1::2]
        c = block[1::2, 0::2]
        d = block[1::2, 1::2]
        ll = (a + b + c + d) / 2.0
        lh = (a - b + c - d) / 2.0
        hl = (a + b - c - d) / 2.0
        hh = (a - b - c + d) / 2.0
        half = size // 2
        x[:half, :half] = ll
        x[:half, half:size] = lh
        x[half:size, :half] = hl
        x[half:size, half:size] = hh
        size = half
    return x


def inverse_haar2d_single(coeff: np.ndarray, levels: int = 3) -> np.ndarray:
    """Inverse of haar2d_single."""
    x = np.asarray(coeff, dtype=np.float32).copy()
    h, w = x.shape
    max_levels = int(np.log2(h))
    levels = min(levels, max_levels)
    sizes = [2 ** i for i in range(1, levels + 1)]
    for size in sizes:
        half = size // 2
        ll = x[:half, :half].copy()
        lh = x[:half, half:size].copy()
        hl = x[half:size, :half].copy()
        hh = x[half:size, half:size].copy()
        a = (ll + lh + hl + hh) / 2.0
        b = (ll - lh + hl - hh) / 2.0
        c = (ll + lh - hl - hh) / 2.0
        d = (ll - lh - hl + hh) / 2.0
        block = np.zeros((size, size), dtype=np.float32)
        block[0::2, 0::2] = a
        block[0::2, 1::2] = b
        block[1::2, 0::2] = c
        block[1::2, 1::2] = d
        x[:size, :size] = block
    return x


def haar2d_batch(X: np.ndarray, levels: int = 3) -> np.ndarray:
    """Batch Haar wavelet features."""
    imgs = _as_images(X)
    coeffs = np.stack([haar2d_single(img, levels=levels).reshape(-1) for img in imgs], axis=0)
    return coeffs.astype(np.float32)


def scattering_like_features(X: np.ndarray, levels: int = 3) -> np.ndarray:
    """A compact deterministic scattering-like feature map.

    True scattering networks require a bank of complex wavelets. For a small
    offline repo, we approximate the paper's wavelet/scattering idea by using
    Haar coefficients, coefficient magnitudes, and energy coefficients. This
    gives deformation-stable, modulus-style spectral features while remaining
    dependency-free.
    """
    c = haar2d_batch(X, levels=levels)
    return np.concatenate([c, np.abs(c), c * c], axis=1).astype(np.float32)


def torch_haar2d(x: torch.Tensor, levels: int = 3) -> torch.Tensor:
    """Differentiable 2D Haar transform for tensors shaped (B, 1, 8, 8).

    This version avoids in-place writes so autograd remains fast and stable.
    It returns coefficients flattened in multiresolution order rather than
    quadrant-packed image order.
    """
    if x.ndim != 4 or x.shape[1] != 1:
        raise ValueError("Expected tensor of shape (B,1,H,W)")
    current = x
    details = []
    for _ in range(levels):
        if current.shape[-1] < 2 or current.shape[-2] < 2:
            break
        a = current[:, :, 0::2, 0::2]
        b = current[:, :, 0::2, 1::2]
        c = current[:, :, 1::2, 0::2]
        d = current[:, :, 1::2, 1::2]
        ll = (a + b + c + d) / 2.0
        lh = (a - b + c - d) / 2.0
        hl = (a + b - c - d) / 2.0
        hh = (a - b - c + d) / 2.0
        details.extend([lh.flatten(1), hl.flatten(1), hh.flatten(1)])
        current = ll
    details.append(current.flatten(1))
    return torch.cat(details, dim=1)


def torch_wavelet_features(x: torch.Tensor, levels: int = 3) -> torch.Tensor:
    """Flattened differentiable Haar features for learnable spectral operator."""
    return torch_haar2d(x, levels=levels)