# ==============================================================================
# Vision Transformer Interpretability Service
# Functions below are directly adapted line-for-line from:
# draft/f2hf_ablation_study.py (ICCIT paper codebase)
# ==============================================================================

import io
import json
import os
import math
import warnings
import base64
import collections.abc
from itertools import repeat
from functools import partial
from typing import Dict, List, Optional, Tuple, Any

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import torchvision.transforms as transforms

# Optional einops with native PyTorch fallback
try:
    from einops import rearrange
except ImportError:
    def rearrange(tensor, pattern, **kwargs):
        if pattern == 'b n (qkv h d) -> qkv b h n d':
            h = kwargs['h']
            qkv = kwargs['qkv']
            b, n, total_dim = tensor.shape
            d = total_dim // (qkv * h)
            return tensor.reshape(b, n, qkv, h, d).permute(2, 0, 3, 1, 4)
        elif pattern == 'b h n d -> b n (h d)':
            b, h, n, d = tensor.shape
            return tensor.permute(0, 2, 1, 3).reshape(b, n, h * d)
        raise NotImplementedError(f"Pattern {pattern} not implemented in fallback")


# ------------------------------------------------------------------------------
# Lines 467-476 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            return x
        return tuple(repeat(x, n))
    return parse

to_2tuple = _ntuple(2)


# ------------------------------------------------------------------------------
# Lines 622-676 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    def norm_cdf(x):
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)
        tensor.uniform_(2 * l - 1, 2 * u - 1)
        tensor.erfinv_()
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


# ------------------------------------------------------------------------------
# Lines 322-329 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
def _conv_filter(state_dict, patch_size=16):
    """ convert patch embedding weight from manual patchify + linear proj to conv"""
    out_dict = {}
    for k, v in state_dict.items():
        if 'patch_embed.proj.weight' in k:
            if len(v.shape) == 2:
                v = v.reshape((v.shape[0], 3, patch_size, patch_size))
        out_dict[k] = v
    return out_dict


# ------------------------------------------------------------------------------
# Lines 680-701 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x


# ------------------------------------------------------------------------------
# Lines 705-746 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def save_fc1_gradient(self, fc1_grad):
        self.fc1_gradient = fc1_grad

    def get_fc1_gradient(self):
        return self.fc1_gradient

    def save_fc2_gradient(self, fc2_grad):
        self.fc2_gradient = fc2_grad

    def get_fc2_gradient(self):
        return self.fc2_gradient

    def save_act_gradient(self, act_grad):
        self.act_gradient = act_grad

    def get_act_gradient(self):
        return self.act_gradient

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


# ------------------------------------------------------------------------------
# Lines 750-835 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., verbose=False):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.attn_gradients = None
        self.attention_map = None
        self.vproj = None
        self.input = None

    def save_input(self, z):
        self.input = z

    def get_input(self):
        return self.input

    def save_attn_gradients(self, attn_gradients):
        self.attn_gradients = attn_gradients

    def get_attn_gradients(self):
        return self.attn_gradients

    def save_attention_map(self, attention_map):
        self.attention_map = attention_map

    def get_attention_map(self):
        return self.attention_map

    def forward(self, x, register_hook=False):
        b, n, _ = x.shape
        h = self.num_heads
        self.save_input(x)

        qkv = self.qkv(x)
        q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h)

        W_qkv = self.qkv.weight
        dim = x.shape[-1]

        W_v = W_qkv[2 * dim : 3 * dim, :]
        self.W_v = W_v
        self.W_proj = self.proj.weight

        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
        attn = dots.softmax(dim=-1)
        attn = self.attn_drop(attn)
        attention_score = torch.einsum('bhij,bhjd->bhid', attn, v)

        self.save_attention_map(attn)
        if register_hook:
            attn.register_hook(self.save_attn_gradients)

        attention_score_rearranged = rearrange(attention_score, 'b h n d -> b n (h d)')
        out = self.proj(attention_score_rearranged)
        out = self.proj_drop(out)

        self.vproj = torch.matmul(rearrange(v, 'b h n d -> b n (h d)'), self.proj.weight.t())
        return out


# ------------------------------------------------------------------------------
# Lines 840-892 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, block_index=None, verbose=False):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, verbose=verbose)
        self.norm2 = norm_layer(dim)

        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.input = None
        self.output = None
        self.tilde = None
        self.block_index = block_index
        self.verbose = verbose

    def save_input(self, z):
        self.input = z

    def get_input(self):
        return self.input

    def save_tilde(self, z):
        self.tilde = z

    def get_tilde(self):
        return self.tilde

    def save_output(self, z):
        self.output = z

    def get_output(self):
        return self.output

    def forward(self, x, register_hook=False):
        self.save_input(x)
        attn_out = self.attn(self.norm1(x), register_hook=register_hook)
        out1 = x + attn_out
        self.save_tilde(out1)

        mlp_out = self.mlp(self.norm2(out1))
        out = out1 + mlp_out
        self.save_output(out)
        return out


# ------------------------------------------------------------------------------
# Lines 896-997 of draft/f2hf_ablation_study.py
# ------------------------------------------------------------------------------
class VisionTransformer(nn.Module):
    """ Vision Transformer
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=False, drop_rate=0., attn_drop_rate=0., norm_layer=nn.LayerNorm,
                 global_pooling=False, dino=False, verbose=False):
        super().__init__()

        self.global_pooling = global_pooling
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
                drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=norm_layer, block_index=i, verbose=verbose)
            for i in range(depth)])
        if global_pooling:
            self.fc_norm = norm_layer(embed_dim)
        else:
            self.norm = norm_layer(embed_dim)

        if dino:
            self.head = nn.Linear(embed_dim * 2, num_classes) if num_classes > 0 else nn.Identity()
        else:
            self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        self.cls_gradients = None
        self.input_grad = None
        self.dino = dino

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, x, register_hook=False):
        B = x.shape[0]
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        for blk in self.blocks[:-1]:
            x = blk(x, register_hook=register_hook)

        x = self.blocks[-1](x, register_hook=register_hook)

        if self.global_pooling:
            x = x[:, 1:, :].mean(axis=1)
            x = self.fc_norm(x)
        elif self.dino:
            x = self.norm(x)
            x = torch.cat((x[:, 0].unsqueeze(-1), x[:, 1:, :].mean(axis=1).unsqueeze(-1)), dim=-1)
            x = x.reshape(x.shape[0], -1)
        else:
            x = self.norm(x)
            x = x[:, 0]

        x = self.head(x)
        return x


# ------------------------------------------------------------------------------
# Lines 1145-1187 of draft/f2hf_ablation_study.py (Attention Rollout)
# ------------------------------------------------------------------------------
def compute_rollout_attention(all_layer_matrices, start_layer=0):
    num_tokens = all_layer_matrices[0].shape[1]
    batch_size = all_layer_matrices[0].shape[0]
    eye = torch.eye(num_tokens).expand(batch_size, num_tokens, num_tokens).to(all_layer_matrices[0].device)
    all_layer_matrices = [all_layer_matrices[i] + eye for i in range(len(all_layer_matrices))]
    matrices_aug = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True)
                          for i in range(len(all_layer_matrices))]
    joint_attention = matrices_aug[start_layer]
    for i in range(start_layer+1, len(matrices_aug)):
        joint_attention = matrices_aug[i].bmm(joint_attention)
    return joint_attention


def generate_rollout(model, input_tensor, device=None, start_layer=0, mae=False):
    input_tensor = input_tensor.to(device)
    out = model(input_tensor)

    all_layer_attentions = []
    for blk in model.blocks:
        attn = blk.attn.get_attention_map()
        attn_avg = attn.mean(dim=1).detach()
        all_layer_attentions.append(attn_avg)

    rollout = compute_rollout_attention(all_layer_attentions, start_layer=start_layer)
    if mae:
        return rollout[:, 1:, 1:].mean(1)
    else:
        return rollout[:, 0, 1:]


# ------------------------------------------------------------------------------
# Lines 1050-1141 of draft/f2hf_ablation_study.py (F2HF & BeyondIntuition)
# ------------------------------------------------------------------------------
def beyond_intuition_tokenwise(model, x, device, onehot=None, ffd_consideration=False, index=None, steps=20, start_layer=1, samples=20, noise=0.1, mae=False, ssl=False, dino=False, taken_head_idx = None):
    b = x.shape[0]
    batch_size = b
    x = x.to(device)

    output = model(x)
    logit_score = output

    if index is None:
        index = np.argmax(output.cpu().data.numpy(), axis=-1)
    elif isinstance(index, int):
        index = np.array([index])

    out_sum = torch.sum(output)

    one_hot = np.zeros((batch_size, output.size()[-1]), dtype=np.float32)
    one_hot[np.arange(batch_size), index] = 1
    one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device)
    one_hot = torch.sum(one_hot)

    model.zero_grad()
    one_hot.backward(retain_graph=True)

    _, num_head, num_tokens, _ = model.blocks[-1].attn.get_attention_map().shape
    R = torch.eye(num_tokens, num_tokens).expand(b, num_tokens, num_tokens).to(device)

    for blk_idx, blk in enumerate(model.blocks):
        if blk_idx < start_layer - 1:
            continue

        # Calculate alpha in the paper
        fc2_w = blk.mlp.fc2.weight
        h_dim = fc2_w.shape[1] // num_head
        impw = torch.abs(fc2_w.reshape((num_head, h_dim, -1))).mean(2).mean(1).unsqueeze(1).unsqueeze(2)

        z = blk.get_input()
        vproj = blk.attn.vproj
        order = torch.linalg.norm(vproj, dim=-1).squeeze() / (torch.linalg.norm(z, dim=-1).squeeze() + 1e-8)
        m = torch.diag_embed(order)

        cam = blk.attn.get_attention_map()
        cam = cam.reshape(-1, cam.shape[-2], cam.shape[-1])

        if taken_head_idx is None:
            if ffd_consideration:
                cam = cam * impw
            cam = cam.mean(0)
        else:
            cam = cam[taken_head_idx]

        O_t = torch.matmul(cam.to(device), m.to(device))
        R = R + torch.matmul(O_t, R.to(device))

    if ssl:
        if mae:
            return R[:, 1:, 1:].abs().mean(axis=1), logit_score
        elif dino:
            return R[:, 1:, 1:].abs().mean(axis=1) + R[:, 0, 1:].abs(), logit_score
        else:
            return R[:, 0, 1:].abs(), logit_score

    total_gradients = torch.zeros(b, num_head, num_tokens, num_tokens, device=device)
    for alpha in np.linspace(0, 1, steps):
        data_scaled = x * alpha
        output = model(data_scaled, register_hook=True)
        one_hot = np.zeros((b, output.size()[-1]), dtype=np.float32)
        one_hot[np.arange(b), index] = 1
        one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device)
        one_hot = torch.sum(one_hot * output)

        model.zero_grad()
        one_hot.backward(retain_graph=True)

        gradients = model.blocks[-1].attn.get_attn_gradients()
        if gradients is not None:
            total_gradients += gradients

    W_state = (total_gradients / steps).clamp(min=0).mean(1).reshape(b, num_tokens, num_tokens)
    R = W_state * R.abs()

    if mae:
        R_MAE = R[:, 1:, 1:].mean(axis=1)
        return (R_MAE, logit_score)
    elif dino:
        R_DINO = (R[:, 1:, 1:].mean(axis=1) + R[:, 0, 1:])
        return (R_DINO, logit_score)
    else:
        R_general = R[:, 0, 1:]
        return (R_general, logit_score)


# ------------------------------------------------------------------------------
# Lines 1316-1364 of draft/f2hf_ablation_study.py (Integrated Gradients)
# ------------------------------------------------------------------------------
def integrated_gradient(model, x, device, index=None, steps=20, mae=False, dino=False, ssl=False):
    b = x.shape[0]
    batch_size = b
    x = x.to(device)

    output = model(x)
    if index is None:
        index = np.argmax(output.cpu().data.numpy(), axis=-1)
    elif isinstance(index, int):
        index = np.array([index])

    one_hot = np.zeros((batch_size, output.size()[-1]), dtype=np.float32)
    one_hot[np.arange(batch_size), index] = 1
    one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device)
    one_hot = torch.sum(one_hot)

    model.zero_grad()
    one_hot.backward(retain_graph=True)

    _, num_head, num_tokens, _ = model.blocks[-1].attn.get_attention_map().shape
    total_gradients = torch.zeros(b, num_head, num_tokens, num_tokens, device=device)

    for alpha in np.linspace(0, 1, steps):
        data_scaled = x * alpha
        output = model(data_scaled, register_hook=True)
        one_hot = np.zeros((b, output.size()[-1]), dtype=np.float32)
        one_hot[np.arange(b), index] = 1
        one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device)
        one_hot = torch.sum(one_hot * output)

        model.zero_grad()
        one_hot.backward(retain_graph=True)

        gradients = model.blocks[-1].attn.get_attn_gradients()
        if gradients is not None:
            total_gradients += gradients

    W_state = (total_gradients / steps).clamp(min=0).mean(1).reshape(b, num_tokens, num_tokens)
    R = W_state

    if mae:
        R_MAE = R[:, 1:, 1:].mean(axis=1)
        return R_MAE
    elif dino:
        R_DINO = (R[:, 1:, 1:].mean(axis=1) + R[:, 0, 1:])
        return R_DINO
    else:
        R_general = R[:, 0, 1:]
        return R_general


# ------------------------------------------------------------------------------
# Lines 1455-1487 of draft/f2hf_ablation_study.py (Grad-CAM)
# ------------------------------------------------------------------------------
def gradcam(model, x, device, target_cls_idx=None, mae=False):
    x = x.to(device)
    output = model(x, register_hook=True)

    if target_cls_idx is None:
        target_cls_idx = np.argmax(output.cpu().data.numpy())

    one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
    one_hot[0][target_cls_idx] = 1
    one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device)
    one_hot = torch.sum(one_hot * output)

    model.zero_grad()
    one_hot.backward(retain_graph=True)

    grad = model.blocks[-1].attn.get_attn_gradients().to(device)
    cam = model.blocks[-1].attn.get_attention_map().to(device)

    if mae:
        cam = cam[0, :, 1:, 1:]
        grad = grad[0, :, 1:, 1:]
        grad = grad.mean(dim=[0, 2], keepdim=True)
        cam = (cam * grad).mean(0).mean(1).clamp(min=0)
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    else:
        cam = cam[0, :, 0, 1:].reshape(-1, 14, 14)
        grad = grad[0, :, 0, 1:].reshape(-1, 14, 14)
        grad = grad.mean(dim=[1, 2], keepdim=True)
        cam = (cam * grad).mean(0).clamp(min=0)
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

    return cam


def normalize_class_name(s: str) -> str:
    import re
    # Remove WordNet synset prefix if present (e.g., "n01818515_ruffed_grouse" -> "ruffed_grouse")
    s = re.sub(r'^n\d{8}[_\s]+', '', str(s).strip().lower())
    # Replace underscores, hyphens, and multiple spaces with a single space
    return re.sub(r'[\s_\-]+', ' ', s).strip()


def compact_class_name(s: str) -> str:
    import re
    return re.sub(r'[^a-z0-9]', '', str(s).lower())


# ==============================================================================
# Service Wrapper for Local Pretrained Weights & Web API
# ==============================================================================
class ViTExplanationService:
    def __init__(self, base_dir: str = "."):
        self.base_dir = base_dir
        # In Hugging Face ZeroGPU, PyTorch CUDA is mocked. Run inference on CPU to avoid ZeroGPU token intercept.
        if os.environ.get("SPACES_ZERO_GPU") or os.environ.get("FORCE_CPU"):
            self.device = torch.device("cpu")
        else:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.models: Dict[str, VisionTransformer] = {}
        self.id2label: Dict[int, str] = {}
        self.label2id: Dict[str, int] = {}
        self.norm_label2id: Dict[str, int] = {}
        self.compact_label2id: Dict[str, int] = {}
        self._load_classes()

        # Image transform from line 1023-1027 of draft/f2hf_ablation_study.py
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])

    def _index_label(self, idx: int, label_str: str):
        self.id2label[idx] = label_str
        candidates = [label_str] + [p.strip() for p in label_str.split(",") if p.strip()]
        for cand in candidates:
            raw_l = cand.lower()
            self.label2id[raw_l] = idx
            self.label2id[cand.replace(" ", "_").lower()] = idx
            norm = normalize_class_name(cand)
            if norm:
                self.norm_label2id[norm] = idx
            comp = compact_class_name(cand)
            if comp:
                self.compact_label2id[comp] = idx

    def _load_classes(self):
        possible_paths = [
            os.path.join(self.base_dir, "draft", "vit_config.json"),
            os.path.join(self.base_dir, "vit_config.json"),
            os.path.join(self.base_dir, "static", "classes.json")
        ]
        loaded = False
        for p in possible_paths:
            if os.path.exists(p):
                try:
                    with open(p, "r") as f:
                        data = json.load(f)
                        if "id2label" in data:
                            for k, v in data["id2label"].items():
                                self._index_label(int(k), v)
                            loaded = True
                            break
                        elif isinstance(data, list):
                            for item in data:
                                self._index_label(int(item["index"]), item["name"])
                            loaded = True
                            break
                except Exception as e:
                    print(f"Notice: error reading {p}: {e}")

        if not loaded:
            for i in range(1000):
                self._index_label(i, f"Class {i}")

    def get_classes_list(self) -> List[Dict[str, Any]]:
        return [{"index": i, "name": self.id2label.get(i, f"Class {i}")} for i in range(1000)]

    def resolve_class_index(self, query: str, query_type: str = "auto") -> Tuple[int, str]:
        query_clean = str(query).strip()
        if not query_clean:
            return 0, self.id2label.get(0, "tench")

        if query_type == "number" or (query_type == "auto" and query_clean.isdigit()):
            idx = int(query_clean)
            idx = max(0, min(999, idx))
            return idx, self.id2label.get(idx, f"Class {idx}")

        q_lower = query_clean.lower()
        q_norm = normalize_class_name(query_clean)
        q_comp = compact_class_name(query_clean)

        # 1. Exact match in raw, normalized, or compact lookup
        if q_lower in self.label2id:
            idx = self.label2id[q_lower]
            return idx, self.id2label[idx]
        if q_norm in self.norm_label2id:
            idx = self.norm_label2id[q_norm]
            return idx, self.id2label[idx]
        if q_comp in self.compact_label2id:
            idx = self.compact_label2id[q_comp]
            return idx, self.id2label[idx]

        # 2. Substring matching across normalized dictionary
        for name, idx in self.norm_label2id.items():
            if q_norm in name or name in q_norm:
                return idx, self.id2label[idx]

        # 3. Substring matching across compact dictionary
        for comp, idx in self.compact_label2id.items():
            if q_comp in comp or comp in q_comp:
                return idx, self.id2label[idx]

        return 0, self.id2label.get(0, "tench")

    def get_model(self, model_key: str = "vit_base_patch16_224") -> VisionTransformer:
        if model_key in self.models:
            return self.models[model_key]

        models_dir = os.path.join(self.base_dir, "models")
        if model_key == "vit_small_patch16_224":
            model = VisionTransformer(
                patch_size=16, embed_dim=768, depth=8, num_classes=1000, num_heads=6,
                mlp_ratio=3, qkv_bias=False, norm_layer=partial(nn.LayerNorm, eps=1e-6)
            )
            ckpt_path = os.path.join(models_dir, "vit_small_p16_224-15ec54c9.pth")
        else:
            model = VisionTransformer(
                patch_size=16, embed_dim=768, depth=12, num_classes=1000, num_heads=12,
                mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6)
            )
            ckpt_path = os.path.join(models_dir, "jx_vit_base_p16_224-80ecf9dd.pth")

        if os.path.exists(ckpt_path):
            state = torch.load(ckpt_path, map_location="cpu")
            state = _conv_filter(state)
            model.load_state_dict(state, strict=False)

        model.to(self.device)
        model.eval()
        self.models[model_key] = model
        return model

    def _render_heatmaps(self, raw_map: torch.Tensor, original_pil: Image.Image) -> Tuple[str, str]:
        arr = raw_map.detach().cpu().numpy()
        if arr.size == 196:
            arr = arr.reshape(14, 14)
        elif arr.shape[-2:] != (14, 14):
            arr = arr.reshape(14, 14)

        # 1. Raw saliency map with Reds colormap (matching draft/f2hf_ablation_study.py)
        # Directly streamed from an in-memory RAM buffer — never written to disk
        fig, ax = plt.subplots(figsize=(3, 3))
        ax.imshow(arr, cmap="Reds")
        ax.axis("off")
        fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
        raw_buf = io.BytesIO()
        fig.savefig(raw_buf, format="png", bbox_inches="tight", pad_inches=0)
        plt.close(fig)
        saliency_base64 = "data:image/png;base64," + base64.b64encode(raw_buf.getvalue()).decode("utf-8")

        # 2. Superimposed heatmap on top of original image (seismic overlay)
        h_min, h_max = float(arr.min()), float(arr.max())
        if h_max > h_min:
            norm_map = (arr - h_min) / (h_max - h_min)
        else:
            norm_map = np.zeros_like(arr)

        norm_tensor = torch.from_numpy(norm_map).unsqueeze(0).unsqueeze(0).float()
        upsampled = F.interpolate(norm_tensor, size=(224, 224), mode='bilinear', align_corners=False).squeeze().numpy()
        seismic_rgba = cm.seismic(upsampled)
        seismic_rgb = (seismic_rgba[:, :, :3] * 255).astype(np.uint8)
        orig_resized = original_pil.resize((224, 224)).convert("RGB")
        orig_arr = np.array(orig_resized)
        blended = (0.55 * orig_arr + 0.45 * seismic_rgb).astype(np.uint8)
        overlay_buf = io.BytesIO()
        Image.fromarray(blended).save(overlay_buf, format="PNG")
        overlay_base64 = "data:image/png;base64," + base64.b64encode(overlay_buf.getvalue()).decode("utf-8")

        return saliency_base64, overlay_base64

    def run_simulation(
        self,
        pil_image: Image.Image,
        class_query: str,
        query_type: str = "auto",
        model_name: str = "vit_base_patch16_224",
        selected_methods: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        if selected_methods is None or len(selected_methods) == 0:
            selected_methods = ["f2hf", "beyond_intuition", "integrated_gradients", "gradcam", "rollout"]

        model = self.get_model(model_name)
        input_tensor = self.transform(pil_image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            logits = model(input_tensor)
            probs = F.softmax(logits, dim=-1).squeeze().cpu().numpy()

        top5_indices = np.argsort(probs)[::-1][:5]
        top5 = [
            {"index": int(i), "name": self.id2label.get(int(i), f"Class {i}"), "score": float(probs[i])}
            for i in top5_indices
        ]
        pred_top1 = top5[0]

        target_idx, target_name = self.resolve_class_index(class_query, query_type)
        heatmaps: Dict[str, Any] = {}

        # 1. F2HF: beyond_intuition_tokenwise with ffd_consideration=True
        if "f2hf" in selected_methods:
            h_f2hf, _ = beyond_intuition_tokenwise(model, input_tensor, self.device, index=target_idx, ffd_consideration=True)
            sal_b64, over_b64 = self._render_heatmaps(h_f2hf, pil_image)
            heatmaps["f2hf"] = {
                "id": "f2hf",
                "name": "F2HF",
                "badge": "Proposed Method",
                "description": "Feed-Forward Network Filtered Head Attention",
                "image": sal_b64,
                "saliency_image": sal_b64,
                "overlay_image": over_b64
            }

        # 2. BeyondIntuition: beyond_intuition_tokenwise with ffd_consideration=False, start_layer=1
        if "beyond_intuition" in selected_methods:
            h_bi, _ = beyond_intuition_tokenwise(model, input_tensor, self.device, index=target_idx, ffd_consideration=False, start_layer=1)
            sal_b64, over_b64 = self._render_heatmaps(h_bi, pil_image)
            heatmaps["beyond_intuition"] = {
                "id": "beyond_intuition",
                "name": "BeyondIntuition",
                "badge": "Baseline",
                "description": "Tokenwise Self-Attention Attribution",
                "image": sal_b64,
                "saliency_image": sal_b64,
                "overlay_image": over_b64
            }

        # 3. Integrated Gradients
        if "integrated_gradients" in selected_methods:
            h_ig = integrated_gradient(model, input_tensor, self.device, index=target_idx, steps=10)
            sal_b64, over_b64 = self._render_heatmaps(h_ig, pil_image)
            heatmaps["integrated_gradients"] = {
                "id": "integrated_gradients",
                "name": "Integrated Gradients",
                "badge": "Baseline",
                "description": "Path-Integrated Head Attentions",
                "image": sal_b64,
                "saliency_image": sal_b64,
                "overlay_image": over_b64
            }

        # 4. Grad-CAM
        if "gradcam" in selected_methods:
            h_gcam = gradcam(model, input_tensor, self.device, target_cls_idx=target_idx)
            sal_b64, over_b64 = self._render_heatmaps(h_gcam, pil_image)
            heatmaps["gradcam"] = {
                "id": "gradcam",
                "name": "Grad-CAM",
                "badge": "Baseline",
                "description": "Attention Map Weighted by Target Gradients",
                "image": sal_b64,
                "saliency_image": sal_b64,
                "overlay_image": over_b64
            }

        # 5. Attention Rollout
        if "rollout" in selected_methods:
            h_roll = generate_rollout(model, input_tensor, device=self.device, start_layer=1)
            sal_b64, over_b64 = self._render_heatmaps(h_roll, pil_image)
            heatmaps["rollout"] = {
                "id": "rollout",
                "name": "Attention Rollout",
                "badge": "Baseline",
                "description": "Joint Inter-Layer Attention Flow",
                "image": sal_b64,
                "saliency_image": sal_b64,
                "overlay_image": over_b64
            }

        thumb = pil_image.copy()
        thumb.thumbnail((300, 300))
        orig_buf = io.BytesIO()
        thumb.save(orig_buf, format="PNG")
        orig_base64 = "data:image/png;base64," + base64.b64encode(orig_buf.getvalue()).decode("utf-8")

        return {
            "prediction": pred_top1,
            "top5": top5,
            "target": {"index": target_idx, "name": target_name},
            "model_name": model_name,
            "original_image": orig_base64,
            "heatmaps": heatmaps
        }