"""Per-encoder class feature centroids for feature-tower decoy steering (M4).

Feature-only towers (DINOv2, AIMv2, the modern-LLM vision towers) have no text
tower, so in v0.1 they could only contribute an *untargeted* repel term (push the
image off its clean feature). That moves off-manifold but doesn't aim anywhere.

M4 gives each feature tower an explicit target: the mean image feature of a small
pool of exemplar images for the truth class and the decoy class. The attack can
then *steer* the feature toward the decoy centroid and away from the truth centroid
— the same targeted objective contrastive encoders get from their text tower, but
grounded in image features instead of text.

Centroids are computed once per run over a leakage-free exemplar pool (images that
are NOT in the attacked set), keyed by class word so the loss can look up
`centroids[enc.name][decoy]` / `[truth]` directly.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from PIL import Image


def _to_tensor(img: Image.Image, device: str) -> torch.Tensor:
    a = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(a).permute(2, 0, 1).unsqueeze(0).to(device)


def _class_dirs(exemplar_root: Path) -> dict[str, Path]:
    """Map class word -> directory. Layout: exemplar_root/<class word>/*.jpg|png."""
    out: dict[str, Path] = {}
    for d in sorted(exemplar_root.iterdir()):
        if d.is_dir():
            out[d.name.replace("_", " ").strip().lower()] = d
    return out


@torch.no_grad()
def compute_centroids(encoders: list, exemplar_root: str | Path,
                      max_per_class: int = 24, device: str = "cuda",
                      log=print) -> dict[str, dict[str, torch.Tensor]]:
    """Return centroids[enc.name][class_word] = L2-normalized mean image feature.

    Computed for every encoder (contrastive ones can use it too, but by default
    only feature towers consume it in the loss). Missing classes are simply absent.
    """
    root = Path(exemplar_root)
    class_dirs = _class_dirs(root)
    if not class_dirs:
        raise FileNotFoundError(f"no class subdirs under {root}")

    # Load exemplar images once (shared across encoders).
    imgs: dict[str, list[Image.Image]] = {}
    for word, d in class_dirs.items():
        files = [p for p in sorted(d.iterdir())
                 if p.suffix.lower() in (".jpg", ".jpeg", ".png")][:max_per_class]
        imgs[word] = [Image.open(p).convert("RGB") for p in files]

    centroids: dict[str, dict[str, torch.Tensor]] = {}
    for enc in encoders:
        per_class: dict[str, torch.Tensor] = {}
        for word, pool in imgs.items():
            if not pool:
                continue
            feats = []
            for im in pool:
                f = enc.image_feat(_to_tensor(im, device)).detach().float()
                feats.append(f)
            mean = torch.cat(feats, dim=0).mean(dim=0, keepdim=True)
            per_class[word] = (mean / (mean.norm(dim=-1, keepdim=True) + 1e-12))
        centroids[enc.name] = per_class
        log(f"[targets] {enc.name}: centroids for {len(per_class)} classes")
    return centroids