Blame for model/LTF_v3_3.py

"""
LTF V3.3: V3 + pool-logit gated fusion.

Motivation:
- Keep V3 outputs (logits_level1 / logits_level2 / logits_fusion) unchanged for training tools.
- Use "logit_pool" signals (attention-weighted instance logits) as a reliability cue to gate fusion
  between level1_emb and level2_emb, while keeping fusion_head input dim = 2*embed_dim.
"""

from typing import Optional, Tuple

import torch
import torch.nn as nn

from model.LTF import LTFCore
from model.LTF_v2 import (
    BagTokenShuffle,
    _clip_vit_resblock_modules,
    _count_clip_vit_depth,
)
from mil_module.mil_core import MILCore


class LTFCoreV3_3(nn.Module):
    """
    V3 core + a tiny gate that uses logit_pool(level1_tiles) and logit_pool(level2_cls-bag).

    - Level1: patch tokens (tiles only) -> MIL per tile -> tile bag embeddings -> query attention -> level1_emb
      Additionally: compute per-tile instance logits and attention weights -> logit_pool across tiles.
    - Level2: CLS tokens (global + tiles) -> MIL (attn pooling) -> bag_embed = level2_emb
      Additionally: compute logit_pool from MILHead's attn * inst_logits.
    - Fusion: g = sigmoid(MLP([l1_pool_logits, l2_pool_logits])) in (0,1),
      fusion_in = [g*level1_emb, (1-g)*level2_emb] -> fusion_head (same as V3 input dim).
    """

    def __init__(
        self,
        embed_dim=1024,
        num_heads=8,
        num_classes=2,
        num_layers=1,
        dropout=0.1,
        tile_heads=2,
        tile_layers=2,
        use_global_image=True,
        use_patch_transformer=True,
        linear_classifier=False,
        mil_heads=8,
        fusion_hidden=512,
        l1_cross_shuffle_p: float = 0.9,
        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
        l1_with_replacement_p: float = 0.5,
        l2_cross_shuffle_p: float = 0.5,
        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
        l2_with_replacement_p: float = 0.5,
        bag_token_shuffle_p: float = 0.5,
        cross_shuffle_min_keep_original: int = 0,
        use_mil_cross_shuffle: bool = True,
        pool_gate_hidden: int = 16,
        use_pool_gate_fusion: bool = True,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.use_global_image = use_global_image
        self.num_classes = num_classes
        self.use_mil_cross_shuffle = bool(use_mil_cross_shuffle)

        self.l1_cross_shuffle_p = float(l1_cross_shuffle_p)
        self.l1_replace_ratio_range = (
            float(l1_replace_ratio_range[0]),
            float(l1_replace_ratio_range[1]),
        )
        self.l1_with_replacement_p = float(l1_with_replacement_p)
        self.l2_cross_shuffle_p = float(l2_cross_shuffle_p)
        self.l2_replace_ratio_range = (
            float(l2_replace_ratio_range[0]),
            float(l2_replace_ratio_range[1]),
        )
        self.l2_with_replacement_p = float(l2_with_replacement_p)
        self.bag_token_shuffle_p = float(bag_token_shuffle_p)
        self.cross_shuffle_min_keep_original = int(cross_shuffle_min_keep_original)

        self.use_pool_gate_fusion = bool(use_pool_gate_fusion)

        # Keep the same core structure and frozen baseline classifier as V3 for compatibility.
        self.baseline_core = LTFCore(
            embed_dim=embed_dim,
            num_heads=num_heads,
            num_classes=num_classes,
            num_layers=num_layers,
            dropout=dropout,
            tile_heads=tile_heads,
            tile_layers=tile_layers,
            use_global_image=use_global_image,
            use_patch_transformer=use_patch_transformer,
            linear_classifier=linear_classifier,
        )
        for p in self.baseline_core.classifier.parameters():
            p.requires_grad = False

        mil_kw = dict(
            attn_hidden=min(256, embed_dim),
            attn_dropout=dropout,
            attn_bag_mode="embed_cls",
        )
        self.mil_level1 = MILCore(
            d_in=embed_dim,
            num_classes=num_classes,
            mil_head_pooling="attn",
            mil_head_kwargs=mil_kw,
            use_cross_shuffle=self.use_mil_cross_shuffle,
            trans=BagTokenShuffle(),
            allow_trans_token_drop=False,
        )
        self.mil_level2 = MILCore(
            d_in=embed_dim,
            num_classes=num_classes,
            mil_head_pooling="attn",
            mil_head_kwargs=mil_kw,
            use_cross_shuffle=self.use_mil_cross_shuffle,
            trans=BagTokenShuffle(),
            allow_trans_token_drop=False,
        )

        # Level-1 aggregation across tiles (query attention) + tile-level instance logits for logit_pool.
        self.level1_query = nn.Parameter(torch.zeros(1, 1, embed_dim))
        mh = mil_heads if embed_dim % mil_heads == 0 else 8
        while embed_dim % mh != 0:
            mh -= 1
        if mh < 1:
            mh = 1
        self.level1_aggregate_attn = nn.MultiheadAttention(
            embed_dim, mh, dropout=dropout, batch_first=True
        )

        # Heads
        self.head_level1 = nn.Linear(embed_dim, num_classes)
        self.head_level2 = nn.Linear(embed_dim, num_classes)
        self.tile_inst_head = nn.Linear(embed_dim, num_classes)

        fusion_in = embed_dim + embed_dim
        self.fusion_head = nn.Sequential(
            nn.Linear(fusion_in, fusion_hidden),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(fusion_hidden, num_classes),
        )

        # Gate uses pooled logits from L1/L2: concat => [B, 2*C]
        gate_in = 2 * num_classes
        self.pool_gate = nn.Sequential(
            nn.Linear(gate_in, int(pool_gate_hidden)),
            nn.ReLU(),
            nn.Linear(int(pool_gate_hidden), 1),
        )

        nn.init.trunc_normal_(self.level1_query, std=0.02)

    @staticmethod
    def _sample_replace_ratio(low: float, high: float, device: torch.device) -> float:
        lo, hi = min(low, high), max(low, high)
        lo = max(0.0, lo)
        hi = min(1.0, hi)
        if hi <= lo:
            return float(lo)
        u = torch.rand((), device=device)
        return float(lo + (hi - lo) * u)

    def _apply_patch_transformer(self, x):
        b2, n2, t2, d2 = x.shape
        xf = x.reshape(b2 * n2, t2, d2)
        if self.baseline_core.patch_transformer is not None:
            xf = self.baseline_core.patch_transformer(xf)
        return xf.reshape(b2, n2, t2, d2)

    @staticmethod
    def _logit_pool_from_mil_head(mil_head_out: dict) -> torch.Tensor:
        """
        Compute logit_pool = sum(attn * inst_logits) over instances.
        Requires pooling='attn' so that 'attn' is not None.
        """
        inst = mil_head_out.get("inst_logits")
        attn = mil_head_out.get("attn")
        if inst is None or attn is None:
            raise RuntimeError("logit_pool requires MILHead outputs: inst_logits and attn.")
        return (attn * inst).sum(dim=1)

    def forward(
        self,
        x,
        labels: Optional[torch.Tensor] = None,
        x_mil: Optional[torch.Tensor] = None,
        return_viz: bool = False,
    ):
        x_feat = self._apply_patch_transformer(x)
        if x_mil is not None:
            if x_mil.shape != x.shape:
                raise ValueError(f"x_mil shape {tuple(x_mil.shape)} != x shape {tuple(x.shape)}")
            x_mil_feat = self._apply_patch_transformer(x_mil)
        else:
            x_mil_feat = x_feat

        b, n, t, d = x_mil_feat.shape
        if n < 2:
            raise ValueError(f"LTFCoreV3_3 expects N>=2 (global + tiles), got N={n}")
        nt = n - 1
        train_shuffle = self.training and labels is not None
        p_bag_shuffle = self.bag_token_shuffle_p if train_shuffle else 0.0

        # --- Level 1: patch tokens (tiles only, no CLS) ---
        patch_only = x_mil_feat[:, 1:, 1:, :].reshape(b * nt, t - 1, d)
        do1 = train_shuffle and float(torch.rand((), device=x.device)) < self.l1_cross_shuffle_p
        rr1 = self._sample_replace_ratio(*self.l1_replace_ratio_range, x.device)
        wr1 = bool(float(torch.rand((), device=x.device)) < self.l1_with_replacement_p)
        y1 = labels.unsqueeze(1).expand(-1, nt).reshape(-1) if labels is not None else None
        mk = self.cross_shuffle_min_keep_original
        mil1 = self.mil_level1(
            patch_only,
            y=y1,
            return_dict=True,
            do_cross_shuffle=do1,
            replace_ratio=rr1,
            with_replacement=wr1,
            detach_inserted=True,
            cross_shuffle_min_keep=mk,
            p_shuffle=p_bag_shuffle,
        )
        be = mil1["mil_head"]["bag_embed"]
        if be is None:
            raise RuntimeError("MIL level1 expected bag_embed (use attn + embed_cls).")
        bag_seq = be.reshape(b, nt, d)  # per-tile embeddings

        # Query-attn aggregate -> level1_emb
        q = self.level1_query.expand(b, -1, -1)
        agg_out, attn_w = self.level1_aggregate_attn(
            q, bag_seq, bag_seq, need_weights=True, average_attn_weights=False
        )
        level1_emb = agg_out.squeeze(1)  # [B, D]

        # logit_pool across tiles using the same attention weights.
        # attn_w: [B, heads, 1, Nt] when average_attn_weights=False
        if attn_w is None:
            raise RuntimeError("Expected attention weights from level1_aggregate_attn.")
        tile_attn = attn_w.mean(dim=1).squeeze(1)  # [B, Nt]
        tile_inst_logits = self.tile_inst_head(bag_seq)  # [B, Nt, C]
        level1_logit_pool = (tile_attn.unsqueeze(-1) * tile_inst_logits).sum(dim=1)  # [B, C]

        # --- Level 2: CLS tokens (global + tiles) ---
        cls_seq = x_mil_feat[:, :, 0, :]  # [B, N, D]
        do2 = train_shuffle and float(torch.rand((), device=x.device)) < self.l2_cross_shuffle_p
        rr2 = self._sample_replace_ratio(*self.l2_replace_ratio_range, x.device)
        wr2 = bool(float(torch.rand((), device=x.device)) < self.l2_with_replacement_p)
        mil2 = self.mil_level2(
            cls_seq,
            y=labels,
            return_dict=True,
            do_cross_shuffle=do2,
            replace_ratio=rr2,
            with_replacement=wr2,
            detach_inserted=True,
            cross_shuffle_min_keep=mk,
            p_shuffle=p_bag_shuffle,
        )
        mil2_head = mil2["mil_head"]
        be2 = mil2_head["bag_embed"]
        if be2 is None:
            raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).")
        level2_emb = be2  # [B, D]
        level2_logit_pool = self._logit_pool_from_mil_head(mil2_head)  # [B, C]

        logits_level1 = self.head_level1(level1_emb)
        logits_level2 = self.head_level2(level2_emb)

        # --- Fusion: pool-logit gate -> keep fusion input dim unchanged ---
        if self.use_pool_gate_fusion:
            gate_in = torch.cat([level1_logit_pool, level2_logit_pool], dim=-1)  # [B, 2C]
            g = torch.sigmoid(self.pool_gate(gate_in))  # [B, 1]
            fusion_in = torch.cat([g * level1_emb, (1.0 - g) * level2_emb], dim=-1)  # [B, 2D]
        else:
            fusion_in = torch.cat([level1_emb, level2_emb], dim=-1)
        logits_fusion = self.fusion_head(fusion_in)

        out = {
            "logits_level1": logits_level1,
            "logits_level2": logits_level2,
            "logits_fusion": logits_fusion,
        }
        if return_viz:
            mil1_head = mil1.get("mil_head", {})
            out["viz"] = {
                # Branch-level raw outputs
                "logits_level1": logits_level1,
                "logits_level2": logits_level2,
                "logits_fusion": logits_fusion,
                # MIL pooled logits and gate inputs
                "level1_logit_pool": level1_logit_pool,
                "level2_logit_pool": level2_logit_pool,
                # Embeddings before branch heads / fusion
                "level1_emb": level1_emb,
                "level2_emb": level2_emb,
                "fusion_in": fusion_in,
                # Attention maps and instance logits
                "tile_attn": tile_attn,
                "tile_inst_logits": tile_inst_logits,
                "level1_patch_attn": mil1_head.get("attn"),
                "level1_patch_inst_logits": mil1_head.get("inst_logits"),
                "level2_attn": mil2_head.get("attn"),
                "level2_inst_logits": mil2_head.get("inst_logits"),
                "num_tiles": torch.full(
                    (b, 1), nt, device=logits_fusion.device, dtype=torch.int32
                ),
                "num_patches_per_tile": torch.full(
                    (b, 1), (t - 1), device=logits_fusion.device, dtype=torch.int32
                ),
                # Scalar gate
                "gate": g if self.use_pool_gate_fusion else torch.full(
                    (b, 1), 0.5, device=logits_fusion.device, dtype=logits_fusion.dtype
                ),
                "use_pool_gate_fusion": torch.tensor(
                    [1 if self.use_pool_gate_fusion else 0],
                    device=logits_fusion.device,
                    dtype=logits_fusion.dtype,
                ),
            }
        return out


class LTFNetV3_3(nn.Module):
    """Frozen CLIP visual + LTFCoreV3_3. Input: (B, N, 3, H, W)."""

    def __init__(
        self,
        clip_visual,
        num_heads=8,
        num_classes=2,
        num_layers=1,
        dropout=0.2,
        tile_heads=2,
        tile_layers=2,
        use_global_image=True,
        use_patch_transformer=True,
        linear_classifier=False,
        mil_heads=8,
        fusion_hidden=512,
        l1_cross_shuffle_p: float = 0.9,
        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
        l1_with_replacement_p: float = 0.5,
        l2_cross_shuffle_p: float = 0.5,
        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
        l2_with_replacement_p: float = 0.5,
        bag_token_shuffle_p: float = 0.5,
        cross_shuffle_min_keep_original: int = 0,
        use_mil_cross_shuffle: bool = True,
        use_clip_intermediate_for_mil: bool = False,
        clip_mil_stop_after_block_1based: int = 12,
        pool_gate_hidden: int = 16,
        use_pool_gate_fusion: bool = True,
    ):
        super().__init__()
        self.visual = clip_visual
        self.dtype = self.visual.conv1.weight.dtype

        for p in self.visual.parameters():
            p.requires_grad = False
        self.visual.eval()

        self.input_resolution = self.visual.input_resolution
        self.embed_dim = self.visual.conv1.weight.shape[0]

        self.use_clip_intermediate_for_mil = bool(use_clip_intermediate_for_mil)
        self.clip_mil_stop_after_block_1based = int(clip_mil_stop_after_block_1based)
        if self.use_clip_intermediate_for_mil:
            n_bl = _count_clip_vit_depth(self.visual)
            stop = self.clip_mil_stop_after_block_1based
            if stop < 1 or stop > n_bl:
                raise ValueError(
                    f"use_clip_intermediate_for_mil requires 1 <= clip_mil_stop_after_block_1based "
                    f"<= {n_bl} (ViT depth), got {stop}"
                )

        self.core = LTFCoreV3_3(
            embed_dim=self.embed_dim,
            num_heads=num_heads,
            num_classes=num_classes,
            num_layers=num_layers,
            dropout=dropout,
            tile_heads=tile_heads,
            tile_layers=tile_layers,
            use_global_image=use_global_image,
            use_patch_transformer=use_patch_transformer,
            linear_classifier=linear_classifier,
            mil_heads=mil_heads,
            fusion_hidden=fusion_hidden,
            l1_cross_shuffle_p=l1_cross_shuffle_p,
            l1_replace_ratio_range=l1_replace_ratio_range,
            l1_with_replacement_p=l1_with_replacement_p,
            l2_cross_shuffle_p=l2_cross_shuffle_p,
            l2_replace_ratio_range=l2_replace_ratio_range,
            l2_with_replacement_p=l2_with_replacement_p,
            bag_token_shuffle_p=bag_token_shuffle_p,
            cross_shuffle_min_keep_original=cross_shuffle_min_keep_original,
            use_mil_cross_shuffle=use_mil_cross_shuffle,
            pool_gate_hidden=pool_gate_hidden,
            use_pool_gate_fusion=use_pool_gate_fusion,
        )

    def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor:
        x = x_bchw.type(self.dtype)
        x = self.visual.conv1(x)
        x = x.reshape(x.shape[0], x.shape[1], -1)
        x = x.permute(0, 2, 1)
        cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(
            x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device
        )
        x = torch.cat([cls, x], dim=1)
        x = x + self.visual.positional_embedding.to(x.dtype)
        x = self.visual.ln_pre(x)
        return x.permute(1, 0, 2)

    def _visual_tokens_from_imgs(self, imgs_bchw: torch.Tensor):
        x = self._stem_to_lnd(imgs_bchw)
        if not self.use_clip_intermediate_for_mil:
            x = self.visual.transformer(x)
            return x.permute(1, 0, 2), None
        blocks = _clip_vit_resblock_modules(self.visual.transformer)
        stop = self.clip_mil_stop_after_block_1based
        for i in range(stop):
            x = blocks[i](x)
        x_mil = x.permute(1, 0, 2).clone()
        for i in range(stop, len(blocks)):
            x = blocks[i](x)
        return x.permute(1, 0, 2), x_mil

    def forward(
        self,
        x: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        chunk_size=256,
        return_viz: bool = False,
    ):
        chunk_size = getattr(self, "eval_chunk_size", chunk_size)
        b, n, c, h, w = x.shape
        imgs = x.reshape(b * n, c, h, w)
        feats_chunks = []
        mil_chunks = []
        with torch.no_grad():
            for start in range(0, imgs.size(0), chunk_size):
                end = min(start + chunk_size, imgs.size(0))
                f_tok, m_tok = self._visual_tokens_from_imgs(imgs[start:end])
                feats_chunks.append(f_tok)
                if m_tok is not None:
                    mil_chunks.append(m_tok)
        feats_token = torch.cat(feats_chunks, dim=0)
        t = feats_token.shape[-2]
        d = feats_token.shape[-1]
        feats = feats_token.reshape(b, n, t, d).float()
        if not mil_chunks:
            return self.core(feats, labels=labels, x_mil=None, return_viz=return_viz)
        mil_token = torch.cat(mil_chunks, dim=0).reshape(b, n, t, d).float()
        return self.core(feats, labels=labels, x_mil=mil_token, return_viz=return_viz)