"""Per-encoder gradient normalization + VMI variance for the ensemble PGD step.

Two transfer levers live here:

1. Per-encoder gradient L2-normalization. Summing scalar losses then taking one
   backward lets whichever encoder emits the largest-magnitude gradient dominate
   the step (empirically the big CLIP ViTs drown out the feature towers). Instead
   we take one backward PER encoder, normalize each input-gradient to unit L2, and
   average. Every architecture then contributes equally regardless of its native
   gradient scale — this is what makes the feature towers actually pull weight.

2. VMI (variance-tuned momentum, Wang & He 2021). Estimates the gradient variance
   in a neighborhood of the current point and adds it to the momentum accumulator,
   which stabilizes the update direction and reliably improves black-box transfer
   over plain MI-FGSM. Cost = `vmi_n` extra ensemble-gradient evaluations per step.
"""

from __future__ import annotations

import random
from typing import Callable

import torch


def encoder_loss(enc, x, truth_txt, decoy_txt, clean_feat_cache, cfg,
                 centroids=None) -> torch.Tensor:
    """Adversarial objective for one encoder (higher = more fooled).

    Targeted steering (pull toward decoy, away from truth):
      - contrastive encoders use their TEXT tower embeddings of the two labels;
      - feature towers use image-feature CENTROIDS of the two classes when a target
        is provided (M4), otherwise contribute only the untargeted repel term.
    All encoders add the untargeted repel term (push off the clean feature).
    """
    f = enc.image_feat(x)
    ff = f.float()
    loss = torch.zeros((), device=x.device)
    if enc.kind == "contrastive":
        tfeat = enc.text_feat([truth_txt, decoy_txt]).detach().float()
        loss = loss + cfg.w_target * ((ff @ tfeat[1:2].T) - (ff @ tfeat[0:1].T)).mean()
    else:
        cts = centroids.get(enc.name) if centroids else None
        c_decoy = cts.get(decoy_txt) if cts else None
        c_truth = cts.get(truth_txt) if cts else None
        if c_decoy is not None and c_truth is not None:
            loss = loss + cfg.w_target * (
                (ff @ c_decoy.float().T) - (ff @ c_truth.float().T)).mean()
    cf = clean_feat_cache.get(enc.name)
    if cf is not None:
        loss = loss + cfg.w_repel * (-(ff @ cf.float().T).mean())
    return loss


def ensemble_grad(delta: torch.Tensor, x0: torch.Tensor, chosen: list,
                  truth: str, decoy: str, clean_feat: dict, cfg,
                  rng: random.Random, eot_fn: Callable,
                  centroids=None) -> torch.Tensor:
    """Aggregated, per-encoder-L2-normalized input gradient (ascend to fool).

    One backward per encoder so each gradient can be normalized independently.
    Returns a gradient tensor shaped like `delta` (float32).
    """
    agg = torch.zeros_like(x0)
    n = 0
    for enc in chosen:
        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 = eot_fn(x_adv, rng)
            total = total + encoder_loss(enc, x_t, truth, decoy, clean_feat, cfg, centroids)
        total = total / cfg.eot_samples
        g = torch.autograd.grad(total, d, retain_graph=False)[0].float()
        g = g / (g.flatten(1).norm(dim=1).view(-1, 1, 1, 1) + 1e-12)  # unit L2 / encoder
        agg = agg + g
        n += 1
    return agg / max(n, 1)


def vmi_variance(delta: torch.Tensor, cur_grad: torch.Tensor, eps: float,
                 cfg, grad_fn: Callable) -> torch.Tensor:
    """Gradient variance in a uniform neighborhood around `delta` (VMI term).

    v = mean_i grad(delta + r_i) - cur_grad,   r_i ~ U[-beta*eps, beta*eps].
    """
    if cfg.vmi_n <= 0:
        return torch.zeros_like(cur_grad)
    beta = cfg.vmi_beta
    acc = torch.zeros_like(cur_grad)
    for _ in range(cfg.vmi_n):
        r = (torch.rand_like(delta) * 2 - 1) * (beta * eps)
        acc = acc + grad_fn((delta + r).clamp(-eps, eps))
    return acc / cfg.vmi_n - cur_grad