"""Learnable spectral modulation model.

This implements the paper's learnable spectral nonlinearity idea using Haar
coefficients, coefficient-wise shrinkage, gain/phase-like scaling, and a final
linear task operator.
"""

from __future__ import annotations

from dataclasses import dataclass
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from .transforms import torch_wavelet_features
from .metrics import nll_from_proba


class LearnableSpectralOperator(nn.Module):
    """Haar coefficients + learnable soft-threshold/gain + task operator."""

    def __init__(self, input_shape=(1, 8, 8), n_classes: int = 10, levels: int = 3):
        super().__init__()
        self.levels = levels
        n_coeff = int(np.prod(input_shape))
        self.log_theta = nn.Parameter(torch.full((n_coeff,), -2.0))
        self.gain = nn.Parameter(torch.ones(n_coeff))
        self.operator = nn.Linear(n_coeff, n_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        coeff = torch_wavelet_features(x, levels=self.levels)
        theta = torch.nn.functional.softplus(self.log_theta) + 1e-6
        shrink = coeff.abs() / (coeff.abs() + theta)
        modulated = self.gain * shrink * coeff
        return self.operator(modulated)

    @property
    def param_count(self) -> int:
        return sum(p.numel() for p in self.parameters())


@dataclass
class TorchModelWrapper:
    model: LearnableSpectralOperator
    device: str = "cpu"

    @property
    def param_count_(self) -> int:
        return self.model.param_count

    def _tensor(self, X):
        X = np.asarray(X, dtype=np.float32)
        if X.ndim == 3:
            X = X[:, None, :, :]
        if X.max() > 1.0:
            X = X / 16.0
        return torch.from_numpy(X).to(self.device)

    def decision_function(self, X):
        self.model.eval()
        with torch.no_grad():
            logits = self.model(self._tensor(X)).cpu().numpy()
        return logits

    def predict_proba(self, X):
        logits = self.decision_function(X)
        logits = logits - logits.max(axis=1, keepdims=True)
        e = np.exp(logits)
        return e / e.sum(axis=1, keepdims=True)

    def predict(self, X):
        return self.decision_function(X).argmax(axis=1)


def _torch_metrics(wrapper: TorchModelWrapper, X, y, prefix: str) -> dict:
    pred = wrapper.predict(X)
    proba = wrapper.predict_proba(X)
    return {
        f"{prefix}_accuracy": float((pred == y).mean()),
        f"{prefix}_macro_f1": float(__import__("sklearn.metrics").metrics.f1_score(y, pred, average="macro")),
        f"{prefix}_nll": nll_from_proba(proba, y),
    }


def train_learnable_spectral(
    X_train,
    y_train,
    X_val,
    y_val,
    epochs: int = 30,
    lr: float = 3e-2,
    weight_decay: float = 1e-4,
    seed: int = 7,
    batch_size: int = 128,
    device: str = "cpu",
):
    torch.manual_seed(seed)
    model = LearnableSpectralOperator().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    loss_fn = nn.CrossEntropyLoss()

    Xtr = np.asarray(X_train, dtype=np.float32)
    Xva = np.asarray(X_val, dtype=np.float32)
    if Xtr.ndim == 3:
        Xtr = Xtr[:, None, :, :]
    if Xva.ndim == 3:
        Xva = Xva[:, None, :, :]
    ytr = np.asarray(y_train, dtype=np.int64)
    ds = TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr))
    loader = DataLoader(ds, batch_size=batch_size, shuffle=True)

    history = []
    for epoch in range(1, epochs + 1):
        model.train()
        total = 0.0
        seen = 0
        for xb, yb in loader:
            xb = xb.to(device)
            yb = yb.to(device)
            opt.zero_grad(set_to_none=True)
            logits = model(xb)
            # mild sparsity prior on spectral gains/shrinkage, matching the paper's regularized spectral learning theme
            loss = loss_fn(logits, yb) + 1e-4 * model.gain.abs().mean()
            loss.backward()
            opt.step()
            total += float(loss.item()) * len(xb)
            seen += len(xb)
        wrapper = TorchModelWrapper(model, device=device)
        val_metrics = _torch_metrics(wrapper, X_val, y_val, "val")
        history.append({"epoch": epoch, "train_loss": total / seen, **val_metrics})
    return TorchModelWrapper(model, device=device), history


def fit_learnable_threshold_operator(X_train, y_train, X_val, y_val, lam: float = 1e-3, thresholds=None):
    """Fast validation-learned spectral soft-threshold operator.

    This dependency-light variant learns a global spectral threshold by validation
    search, then estimates the task operator with Hilbert-Schmidt ridge. It is
    used in the default benchmark so the repo runs quickly on CPU.
    """
    from .transforms import haar2d_batch
    from .operators import HilbertSchmidtOperatorClassifier

    if thresholds is None:
        thresholds = [0.0, 0.005, 0.01, 0.025, 0.05, 0.10, 0.20]

    Ctr = haar2d_batch(X_train)
    Cva = haar2d_batch(X_val)

    def shrink(C, theta):
        return C * (np.abs(C) / (np.abs(C) + theta + 1e-8))

    best = None
    for theta in thresholds:
        Ztr = shrink(Ctr, theta)
        Zva = shrink(Cva, theta)
        clf = HilbertSchmidtOperatorClassifier(ridge_lambda=lam).fit(Ztr, y_train)
        val_acc = float((clf.predict(Zva) == y_val).mean())
        if best is None or val_acc > best[0]:
            best = (val_acc, theta, clf)

    class ThresholdSpectralWrapper:
        theta_ = best[1]
        clf_ = best[2]
        param_count_ = int(best[2].param_count_ + 1)

        def _features(self, X):
            C = haar2d_batch(X)
            return shrink(C, self.theta_)

        def decision_function(self, X):
            return self.clf_.decision_function(self._features(X))

        def predict_proba(self, X):
            return self.clf_.predict_proba(self._features(X))

        def predict(self, X):
            return self.clf_.predict(self._features(X))

    return ThresholdSpectralWrapper(), {"selected_threshold": float(best[1]), "val_accuracy": float(best[0])}