"""Uniform, differentiable encoder interface for the ensemble PGD attack.

Every encoder exposes:
  - image_feat(x): [0,1] (B,3,H,W) -> L2-normalized image embedding (differentiable)
  - text_feat(list[str]) -> L2-normalized text embeddings (contrastive encoders only)
  - kind: "contrastive" (has text tower -> targeted decoy loss) or "feature"
          (raw patch features -> untargeted feature-displacement loss)

The contrastive backbone = open_clip encoders (native text direction, cheap,
architecturally diverse incl. a ConvNeXt CNN). HF VLM vision towers (MoonViT,
InternViT, MiniMax-M3) are added as "feature" encoders via ensemble/towers/.
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass

import torch
import torch.nn.functional as F

_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
_TOWERS_DIR = os.path.join(os.path.dirname(__file__), "towers")


# name : (open_clip_arch, pretrained_tag). Grouped for architectural diversity.
OPENCLIP_SPECS: dict[str, tuple[str, str]] = {
    # OpenAI / LAION classic ViT
    "clip-L14-openai": ("ViT-L-14", "openai"),
    "clip-H14-laion2b": ("ViT-H-14", "laion2b_s32b_b79k"),
    "clip-bigG14-laion2b": ("ViT-bigG-14", "laion2b_s39b_b160k"),
    # EVA (different pretraining/arch)
    "eva02-L14-336": ("EVA02-L-14-336", "merged2b_s6b_b61k"),
    "eva02-E14": ("EVA02-E-14", "laion2b_s4b_b115k"),
    # SigLIP / SigLIP2 (sigmoid loss, SO400m shape)
    "siglip2-L16-256": ("ViT-L-16-SigLIP2-256", "webli"),
    "siglip2-so400m-378": ("ViT-SO400M-14-SigLIP2-378", "webli"),
    "siglip-so400m-384": ("ViT-SO400M-14-SigLIP-384", "webli"),
    # MetaCLIP + DFN (different data curation)
    "metaclip-L14": ("ViT-L-14", "metaclip_fullcc"),
    # DFN weights were trained with QuickGELU — use the canonical -quickgelu arch so the
    # activation matches (plain arch silently loads GELU and corrupts the features).
    "dfn-H14-378": ("ViT-H-14-378-quickgelu", "dfn5b"),
    "dfn-L14": ("ViT-L-14-quickgelu", "dfn2b"),
    # ConvNeXt — a CNN, not a ViT (true architectural diversity)
    "convnext-xxl-laion2b": ("convnext_xxlarge", "laion2b_s34b_b82k_augreg"),
    # v0.2.1 Tier-1 contrastive giants (differentiability-probed on klaus-1)
    "siglip2-giant-384": ("ViT-gopt-16-SigLIP2-384", "webli"),           # largest public SigLIP2
    "metaclip2-H-worldwide-378": ("ViT-H-14-worldwide-378", "metaclip2_worldwide"),  # newest Meta, multiling
}

# Default train/held-out split so transfer is measured on unseen architectures.
DEFAULT_TRAIN = [
    "clip-L14-openai", "clip-H14-laion2b", "clip-bigG14-laion2b",
    "eva02-L14-336", "siglip2-L16-256", "siglip-so400m-384",
    "metaclip-L14", "convnext-xxl-laion2b",
]
DEFAULT_HELDOUT = ["dfn-H14-378", "dfn-L14", "eva02-E14", "siglip2-so400m-378"]

# Feature towers = actual vision encoders inside recent open VLMs, loaded via timm
# (clean forward_features, differentiable, no flash_attn / trust_remote_code pain).
# No text tower -> "feature" kind (untargeted repel loss only). Chosen for being
# real VLM backbones AND architecturally novel vs the CLIP-family train set.
TIMM_FEATURE_SPECS: dict[str, str] = {
    "aimv2-huge-336": "aimv2_huge_patch14_336.apple_pt",       # AIMv2 (GLM-4V family init)
    "dinov2-L-reg": "vit_large_patch14_reg4_dinov2.lvd142m",   # DINOv2 self-supervised
}

# Ablation train set: the 8 contrastive encoders + 2 real VLM feature towers.
ABLATION_TRAIN = DEFAULT_TRAIN + ["aimv2-huge-336", "dinov2-L-reg"]

# Modern-LLM vision towers, loaded via hand-written differentiable loaders under
# ensemble/towers/ (verified grad-flow on the Sparks). "feature" kind -> untargeted
# repel loss only. These are the actual encoders inside recent open VLMs, so they're
# the closest open proxies for frontier vision behavior.
#   name : (loader_module, load_fn, feat_fn)
CUSTOM_SPECS: dict[str, tuple[str, str, str]] = {
    "moonvit-so400m": ("moonvit_loader", "load_moonvit", "image_feat_moonvit"),   # Kimi K2.7
    "internvit-300m-v2_5": ("internvit_loader", "load_internvit", "image_feat_internvit"),  # InternVL3.5
    "minimax-m3-vision": ("minimax_loader", "load_minimax", "image_feat_minimax"),  # MiniMax-M3
    # v0.2.1 Tier-1 held-out judges (never attacked; cross-architecture on purpose):
    #   agglomerative distillation of CLIP+SigLIP2+DINOv2+SAM ...
    "c-radio-v3-h": ("cradio_loader", "load_cradio", "image_feat_cradio"),
    #   ... and a real frontier VLM captioner tower (Qwen3-VL-8B vision).
    "qwen3vl-8b-vision": ("qwen3vl_loader", "load_qwen3vl", "image_feat_qwen3vl"),
}

# Full train set: ablation set (contrastive + AIMv2/DINOv2) + the 3 modern-LLM towers.
# This is the complete ensemble used for the v0.1 results.
FULL_TRAIN = ABLATION_TRAIN + list(CUSTOM_SPECS)

# --- v0.2 split (M1): the held-out set must contain architecture FAMILIES the train
# set never sees, so "transfer" measures cross-architecture generalization rather than
# CLIP-family self-similarity. Feature/VLM towers now sit in held-out too; with M4
# decoy-image centroids we can score a real flip on them (adv closer to the decoy
# centroid than the truth centroid), not just a contrastive text margin.
V02_TRAIN = [
    # contrastive core (8)
    "clip-L14-openai", "clip-H14-laion2b", "clip-bigG14-laion2b",
    "eva02-L14-336", "siglip2-L16-256", "siglip-so400m-384",
    "metaclip-L14", "convnext-xxl-laion2b",
    # feature / VLM towers kept in train
    "aimv2-huge-336",            # SSL feature family (train)
    "moonvit-so400m",            # VLM tower (Kimi)
    "minimax-m3-vision",         # VLM tower (MiniMax-M3)
]
V02_HELDOUT = [
    "dfn-H14-378", "dfn-L14", "siglip2-so400m-378",  # CLIP-family contrastive anchors
    "dinov2-L-reg",              # SSL feature family NOT in train (cross-arch)
    "internvit-300m-v2_5",       # VLM tower NOT in train (cross-arch)
]

# --- v0.2.1 split (ADOPTED): the v0.2 comparison showed moving DINOv2+InternViT OUT of the
# attack (to serve as held-out) COST attack strength vs v0.1 (esp. on gpt-5.5, 20.7%->6.9%).
# Fix: put the FULL v0.1 ensemble back in the ATTACK, and use SEPARATE never-trained encoders
# as cross-arch JUDGES so measurement no longer cannibalizes the attack. This recovers gpt-5.5
# to 20.7% and lifts pooled frontier flip to 38.8% (best of v0.1/v0.2/v0.2.1). See
# research/v0_2_1_findings.md.
#
# The two Tier-1 contrastive giants (SigLIP2-giant, MetaCLIP2-H) were ablation-gated and
# REJECTED — adding them to the attack hurt frontier (-3.0pp BOTH) and train margin
# (-3.1pp, CI excludes 0) by diluting the per-step subset with near-CLIP clones. The adopted
# attack therefore does NOT include them.
V021_ATTACK = [
    # contrastive core (8)
    "clip-L14-openai", "clip-H14-laion2b", "clip-bigG14-laion2b",
    "eva02-L14-336", "siglip2-L16-256", "siglip-so400m-384",
    "metaclip-L14", "convnext-xxl-laion2b",
    # feature / VLM towers (all back in the attack — recover v0.1 strength)
    "aimv2-huge-336", "dinov2-L-reg",
    "moonvit-so400m", "internvit-300m-v2_5", "minimax-m3-vision",
]
# Opt-in ablation arm: attack + the two REJECTED Tier-1 giants (kept only to reproduce the gate).
V021_ATTACK_GIANTS = V021_ATTACK + ["siglip2-giant-384", "metaclip2-H-worldwide-378"]
# Back-compat alias: the gate baseline == the adopted attack.
V021_ATTACK_BASE = V021_ATTACK
# Held-out JUDGES: never in any attack set. Cross-architecture on purpose.
V021_HELDOUT = [
    "dfn-H14-378", "dfn-L14",     # CLIP-family contrastive anchors
    "c-radio-v3-h",               # agglomerative (CLIP+SigLIP2+DINOv2+SAM) — max-diversity judge
    "qwen3vl-8b-vision",          # real frontier VLM captioner tower — the closest open proxy
]

# Architectural family per encoder — used by stratified subset sampling so a
# single PGD step can't be dominated by near-clones (e.g. all OpenAI-style ViTs).
# Group by backbone architecture, not by pretraining data.
FAMILY: dict[str, str] = {
    "clip-L14-openai": "clip_vit", "clip-H14-laion2b": "clip_vit",
    "clip-bigG14-laion2b": "clip_vit", "metaclip-L14": "clip_vit",
    "dfn-H14-378": "clip_vit", "dfn-L14": "clip_vit",
    "eva02-L14-336": "eva", "eva02-E14": "eva",
    "siglip2-L16-256": "siglip", "siglip2-so400m-378": "siglip",
    "siglip-so400m-384": "siglip",
    "convnext-xxl-laion2b": "convnext",
    "aimv2-huge-336": "aimv2", "dinov2-L-reg": "dinov2",
    "moonvit-so400m": "moonvit", "internvit-300m-v2_5": "internvit",
    "minimax-m3-vision": "minimax",
    # v0.2.1 Tier-1
    "siglip2-giant-384": "siglip", "metaclip2-H-worldwide-378": "metaclip2",
    "c-radio-v3-h": "radio", "qwen3vl-8b-vision": "qwen3vl",
}


def family_of(name: str) -> str:
    return FAMILY.get(name, name)


@dataclass
class Encoder:
    name: str
    kind: str          # "contrastive" | "feature"
    res: int
    model: object
    tokenizer: object
    mean: torch.Tensor
    std: torch.Tensor
    dtype: torch.dtype
    backend: str = "openclip"   # "openclip" | "timm" | "custom"
    family: str = ""
    feat_fn: object = None      # custom towers: feat_fn(model, x[0,1]) -> (1, D)

    def _prep(self, x: torch.Tensor) -> torch.Tensor:
        x = F.interpolate(x, size=(self.res, self.res), mode="bicubic",
                          align_corners=False, antialias=True).clamp(0, 1)
        x = (x - self.mean) / self.std
        return x.to(self.dtype)

    def image_feat(self, x: torch.Tensor) -> torch.Tensor:
        if self.backend == "custom":
            f = self.feat_fn(self.model, x)        # does its own resize/normalize
        elif self.backend == "timm":
            f = self.model(self._prep(x))          # num_classes=0 -> pooled features
            if f.dim() == 3:                        # (B, N, D) if pooling disabled
                f = f.mean(dim=1)
        else:
            f = self.model.encode_image(self._prep(x))
        return f / f.norm(dim=-1, keepdim=True)

    def text_feat(self, texts: list[str]) -> torch.Tensor:
        assert self.kind == "contrastive"
        tok = self.tokenizer(texts).to(_DEVICE)
        f = self.model.encode_text(tok)
        return f / f.norm(dim=-1, keepdim=True)


def _mean_std_from_preprocess(preprocess) -> tuple[list, list]:
    for t in getattr(preprocess, "transforms", []):
        if t.__class__.__name__ == "Normalize":
            return list(t.mean), list(t.std)
    # OpenAI CLIP defaults
    return ([0.48145466, 0.4578275, 0.40821073],
            [0.26862954, 0.26130258, 0.27577711])


def load_openclip(name: str, dtype=torch.float16) -> Encoder:
    import open_clip

    arch, tag = OPENCLIP_SPECS[name]
    # OpenAI + MetaCLIP + DFN weights were trained with QuickGELU; force it to match
    # (DFN also uses a -quickgelu arch name, but keep the flag for belt-and-suspenders).
    force_qgelu = tag in ("openai", "metaclip_400m", "metaclip_fullcc", "dfn2b", "dfn5b")
    model, _, preprocess = open_clip.create_model_and_transforms(
        arch, pretrained=tag, force_quick_gelu=force_qgelu)
    model = model.to(_DEVICE, dtype=dtype).eval()
    for p in model.parameters():
        p.requires_grad_(False)
    tok = open_clip.get_tokenizer(arch)
    mean, std = _mean_std_from_preprocess(preprocess)
    res = getattr(model.visual, "image_size", 224)
    res = res[0] if isinstance(res, (tuple, list)) else int(res)
    m = torch.tensor(mean, device=_DEVICE).view(1, 3, 1, 1)
    s = torch.tensor(std, device=_DEVICE).view(1, 3, 1, 1)
    return Encoder(name, "contrastive", res, model, tok, m, s, dtype,
                   family=family_of(name))


def load_timm_feature(name: str, dtype=torch.float16) -> Encoder:
    import timm

    model = timm.create_model(TIMM_FEATURE_SPECS[name], pretrained=True, num_classes=0)
    model = model.to(_DEVICE, dtype=dtype).eval()
    for p in model.parameters():
        p.requires_grad_(False)
    cfg = timm.data.resolve_model_data_config(model)
    res = cfg["input_size"][-1]
    m = torch.tensor(cfg["mean"], device=_DEVICE).view(1, 3, 1, 1)
    s = torch.tensor(cfg["std"], device=_DEVICE).view(1, 3, 1, 1)
    return Encoder(name, "feature", res, model, None, m, s, dtype,
                   backend="timm", family=family_of(name))


def load_custom_tower(name: str, dtype=torch.bfloat16) -> Encoder:
    """Load a modern-LLM vision tower via its hand-written loader in ensemble/towers/."""
    import importlib

    if _TOWERS_DIR not in sys.path:
        sys.path.insert(0, _TOWERS_DIR)
    mod_name, load_fn, feat_fn = CUSTOM_SPECS[name]
    mod = importlib.import_module(mod_name)
    model = getattr(mod, load_fn)(dtype=dtype)
    fn = getattr(mod, feat_fn)
    dummy = torch.zeros(1, device=_DEVICE)
    return Encoder(name, "feature", 0, model, None, dummy, dummy, dtype,
                   backend="custom", family=family_of(name), feat_fn=fn)


def load_encoder(name: str, dtype=torch.float16) -> Encoder:
    if name in OPENCLIP_SPECS:
        return load_openclip(name, dtype=dtype)
    if name in TIMM_FEATURE_SPECS:
        return load_timm_feature(name, dtype=dtype)
    if name in CUSTOM_SPECS:
        return load_custom_tower(name, dtype=torch.bfloat16)
    raise KeyError(f"unknown encoder {name!r}")