"""Self-contained PatchGuard detector for the A-EYE backend (model 54+ family).

Mirrors aeye_next/models/patchguard.py from the detector repo, but imports the
backend's LOCAL zero_shot_v4 (which uses the transformers>=4.5x CLIP layer call
`causal_attention_mask=None` that this venv needs). New file; nothing existing is
modified. The image decision and the 16x16 heatmap come from the same forward.
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from zero_shot_v4 import ZeroShotV4Detector


class PatchGuardDetector(ZeroShotV4Detector):
    def __init__(
        self,
        clip_backbone: str = "clip-vit-l-14",
        clip_layer: int = 13,
        semantic_dim: int = 512,
        forensic_dim: int = 256,
        frequency_dim: int = 192,
        fft_bins: int = 48,
        image_size: int = 224,
        num_classes: int = 2,
        num_sources: int = 2,
        dropout: float = 0.25,
        source_grl_lambda: float = 0.0,
        freeze_clip: bool = True,
        patch_hidden: int = 256,
        patch_topk_frac: float = 0.25,
    ):
        super().__init__(
            clip_backbone=clip_backbone,
            clip_layer=clip_layer,
            semantic_dim=semantic_dim,
            forensic_dim=forensic_dim,
            frequency_dim=frequency_dim,
            fft_bins=fft_bins,
            image_size=image_size,
            num_classes=num_classes,
            num_sources=num_sources,
            dropout=dropout,
            source_grl_lambda=source_grl_lambda,
            freeze_clip=freeze_clip,
        )
        clip_hidden = int(self.clip.config.hidden_size)
        forensic_map_dim = 192
        self.patch_topk_frac = float(patch_topk_frac)
        self.patch_head = nn.Sequential(
            nn.LayerNorm(clip_hidden + forensic_map_dim),
            nn.Linear(clip_hidden + forensic_map_dim, patch_hidden),
            nn.GELU(),
            nn.Dropout(p=dropout * 0.5),
            nn.Linear(patch_hidden, 1),
        )
        nn.init.trunc_normal_(self.patch_head[-1].weight, std=0.02)
        nn.init.constant_(self.patch_head[-1].bias, -2.0)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.gamma_max = nn.Parameter(torch.zeros(1))

    def _clip_tokens(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            vision = self.clip.vision_model
            hidden = vision.embeddings(pixel_values=x)
            hidden = vision.pre_layrnorm(hidden)
            for idx, layer in enumerate(vision.encoder.layers, start=1):
                layer_out = layer(hidden, attention_mask=None, causal_attention_mask=None)
                hidden = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out
                if idx >= self.clip_layer:
                    break
            return hidden[:, 1:]

    def _forensic_spatial(self, raw: torch.Tensor) -> torch.Tensor:
        branch = self.forensic_branch
        low = F.avg_pool2d(raw, kernel_size=5, stride=1, padding=2)
        residual = raw - low
        x = torch.cat([residual, residual.abs()], dim=1)
        x = branch.stem(x)
        x = branch.stage1(x)
        x = branch.stage2(x)
        return branch.stage3(x)

    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        raw = self._to_raw_rgb(x)
        tokens = self._clip_tokens(x).float()
        pooled = self.semantic_pool(tokens)
        semantic = self.semantic_proj(pooled)

        forensic_map = self._forensic_spatial(raw)
        forensic_vec = self.forensic_branch.proj(
            self.forensic_branch.pool(forensic_map).flatten(1)
        )
        frequency = self.frequency_branch(raw)

        features = torch.cat([semantic, forensic_vec, frequency], dim=1)
        logits = self.head(features)

        grid = int(math.sqrt(tokens.shape[1]))
        fmap = F.interpolate(forensic_map.float(), size=(grid, grid), mode="bilinear", align_corners=False)
        fmap_tokens = fmap.flatten(2).transpose(1, 2)
        patch_logits = self.patch_head(torch.cat([tokens, fmap_tokens], dim=-1)).squeeze(-1)

        k = max(1, int(round(patch_logits.shape[1] * self.patch_topk_frac)))
        patch_summary = patch_logits.topk(k, dim=1).values.mean(dim=1)
        patch_peak = patch_logits.max(dim=1).values
        z_img = (
            (logits[:, 1] - logits[:, 0])
            + self.gamma.squeeze() * patch_summary
            + self.gamma_max.squeeze() * patch_peak
        )
        return {
            "logits": logits,
            "z_img": z_img,
            "patch_logits": patch_logits.view(-1, grid, grid),
            "patch_summary": patch_summary,
            "features": features,
        }


# architecture of model 54-59 (patchguard family)
PATCHGUARD_ARCH = dict(
    clip_backbone="clip-vit-l-14",
    clip_layer=13,
    semantic_dim=512,
    forensic_dim=256,
    frequency_dim=192,
    fft_bins=48,
    image_size=224,
    num_classes=2,
    num_sources=20,
    dropout=0.26,
    source_grl_lambda=0.0,
    freeze_clip=True,
    patch_hidden=256,
    patch_topk_frac=0.08,
)