"""EOT-hardened, ensemble, momentum-PGD image attack.

Four transfer/stealth levers are individually toggleable (all off = plain
momentum-PGD, which the defaults reproduce):
  - per-encoder gradient L2-normalization    (loss.ensemble_grad; cfg.grad_norm)
  - stratified family-aware subset sampling   (sampling.stratified_sample)
  - VMI variance-tuned momentum               (loss.vmi_variance;   cfg.vmi_n>0)
  - perceptual budget: LPIPS penalty + DCT low-pass  (perceptual.*)

Plain behavior = grad_norm off, vmi_n=0, lpips_weight=0, dct_keep=1.0,
max_per_family huge, min_feature=0.

Loss per contrastive encoder:
    w_t*[cos(f,e_decoy) - cos(f,e_truth)]  +  w_r*[-cos(f, f_clean)]
Feature towers contribute only the untargeted repel term.
"""

from __future__ import annotations

import random
from dataclasses import dataclass

import torch

from ensemble.eot import apply_eot
from ensemble.loss import encoder_loss, ensemble_grad, vmi_variance
from ensemble.sampling import stratified_sample


@dataclass
class AttackCfg:
    eps: float = 12 / 255       # L-inf budget
    steps: int = 200
    step_size: float = 1.5 / 255
    subset: int = 4             # encoders sampled per step
    eot_samples: int = 1        # EOT draws averaged per step
    w_target: float = 1.0
    w_repel: float = 0.5
    momentum: float = 0.9
    seed: int = 0
    # --- transfer/stealth levers (defaults off = plain momentum-PGD) ---
    grad_norm: bool = False     # per-encoder unit-L2 gradient aggregation
    vmi_n: int = 0              # VMI neighbor samples (0 = plain MI-FGSM)
    vmi_beta: float = 1.5       # VMI neighborhood radius (x eps)
    max_per_family: int = 99    # per-step cap of same-architecture encoders
    min_feature: int = 0        # min feature towers guaranteed per step
    lpips_weight: float = 0.0   # soft LPIPS penalty weight
    lpips_tau: float = 0.0      # hard LPIPS budget (0 = off); project delta down
    dct_keep: float = 1.0       # DCT low-pass keep-fraction (1.0 = off)


def _legacy_grad(delta, x0, chosen, truth, decoy, clean_feat, cfg, rng, centroids=None):
    """Plain aggregation: sum scalar losses, one backward, no per-encoder norm."""
    d = delta.detach().requires_grad_(True)
    total = torch.zeros((), device=x0.device)
    for _ in range(cfg.eot_samples):
        x_adv = (x0 + d).clamp(0, 1)
        x_t = apply_eot(x_adv, rng, strong=True)
        for enc in chosen:
            total = total + encoder_loss(enc, x_t, truth, decoy, clean_feat, cfg, centroids)
    total = total / (len(chosen) * cfg.eot_samples)
    return torch.autograd.grad(total, d)[0].float()


def _perceptual_grad(delta, x0, lpips_fn):
    """Gradient of LPIPS(adv, clean) wrt delta (to be descended)."""
    d = delta.detach().requires_grad_(True)
    adv = (x0 + d).clamp(0, 1)
    dist = lpips_fn.distance(adv, x0)
    return torch.autograd.grad(dist, d)[0].float(), float(dist.item())


def _unit(g):
    return g / (g.flatten(1).norm(dim=1).view(-1, 1, 1, 1) + 1e-12)


@torch.enable_grad()
def attack_image(x0: torch.Tensor, truth: str, decoy: str, encoders: list,
                 cfg: AttackCfg, lpips_fn=None, centroids=None) -> torch.Tensor:
    """x0: (1,3,H,W) in [0,1] on cuda. Returns adversarial image in [0,1].

    centroids (optional, M4): dict[enc.name][class_word] -> feature centroid, used
    to give feature (no-text) towers a targeted decoy/truth steering objective.
    """
    from ensemble.perceptual import dct_lowpass

    rng = random.Random(cfg.seed)
    clean_feat = {}
    with torch.no_grad():
        for enc in encoders:
            clean_feat[enc.name] = enc.image_feat(x0).detach()

    delta = torch.zeros_like(x0)
    grad_mom = torch.zeros_like(x0)
    for _ in range(cfg.steps):
        if cfg.max_per_family < len(encoders) or cfg.min_feature > 0:
            chosen = stratified_sample(encoders, cfg.subset, rng,
                                       cfg.max_per_family, cfg.min_feature)
        else:
            chosen = encoders if len(encoders) <= cfg.subset \
                else rng.sample(encoders, cfg.subset)

        def grad_fn(dl):
            if cfg.grad_norm:
                return ensemble_grad(dl, x0, chosen, truth, decoy, clean_feat,
                                     cfg, rng, lambda x, r: apply_eot(x, r, strong=True),
                                     centroids=centroids)
            return _legacy_grad(dl, x0, chosen, truth, decoy, clean_feat, cfg, rng,
                                centroids=centroids)

        adv_grad = grad_fn(delta)
        if cfg.vmi_n > 0:
            adv_grad = adv_grad + vmi_variance(delta, adv_grad, cfg.eps, cfg, grad_fn)

        combined = _unit(adv_grad)
        if cfg.lpips_weight > 0.0 and lpips_fn is not None:
            lp_grad, _ = _perceptual_grad(delta, x0, lpips_fn)
            combined = combined - cfg.lpips_weight * _unit(lp_grad)

        with torch.no_grad():
            g = combined / (combined.abs().mean() + 1e-12)      # MI-FGSM normalize
            grad_mom = cfg.momentum * grad_mom + g
            delta = (delta + cfg.step_size * grad_mom.sign()).clamp(-cfg.eps, cfg.eps)
            if cfg.dct_keep < 1.0:
                delta = dct_lowpass(delta, cfg.dct_keep).clamp(-cfg.eps, cfg.eps)
            delta = (x0 + delta).clamp(0, 1) - x0               # keep image in range
            if cfg.lpips_tau > 0.0 and lpips_fn is not None:
                cur = float(lpips_fn.distance((x0 + delta).clamp(0, 1), x0).item())
                if cur > cfg.lpips_tau:
                    delta = delta * (cfg.lpips_tau / cur)       # project onto budget
        delta = delta.detach()
    return (x0 + delta).clamp(0, 1).detach()