classifier_model.py

"""
classifier_model.py
-------------------
Arsitektur PERSIS SAMA dengan yang dipakai saat training di notebook
(Hybrid_ViT_EfficientNet_Brain_Disease_FIXED.ipynb), supaya checkpoint
hasil training (hybrid_vit_efficientnet_brain_best.pth) bisa di-load
tanpa error "Missing/Unexpected key(s)".

Arsitektur:
  - EfficientNet-B3 (CNN backbone)        -> feature map lokal (1536 dim)
  - Custom Vision Transformer (6 layer)    -> dibangun dari nol (bukan
    ViT pretrained HuggingFace), beroperasi di atas feature map CNN
  - Cross-Modal Attention Fusion           -> gabungkan fitur CNN + ViT
  - Classifier Head: Linear -> GELU -> BatchNorm1d -> Dropout -> Linear
"""

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

try:
    from torchvision.models import efficientnet_b3, EfficientNet_B3_Weights
    HAS_WEIGHTS = True
except ImportError:
    from torchvision.models import efficientnet_b3
    HAS_WEIGHTS = False

from src.config import NUM_CLASSES, IMG_SIZE

logger = logging.getLogger("brain_pipeline")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  SUB-MODULES (identik dengan notebook training)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels=1536, patch_size=1, embed_dim=768):
        super().__init__()
        self.proj = nn.Conv2d(in_channels, embed_dim,
                               kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)
        x = x.flatten(2).transpose(1, 2)
        return x


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, return_attn: bool = False):
        B, N, C = x.shape
        qkv = (self.qkv(x)
               .reshape(B, N, 3, self.num_heads, self.head_dim)
               .permute(2, 0, 3, 1, 4))
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        if return_attn:
            return x, attn  # attn: [B, num_heads, N, N] -- dipakai buat heatmap
        return x


class TransformerBlock(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn  = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        hidden = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x, return_attn: bool = False):
        if return_attn:
            attn_out, attn_weights = self.attn(self.norm1(x), return_attn=True)
            x = x + attn_out
            x = x + self.mlp(self.norm2(x))
            return x, attn_weights
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class CrossModalAttentionFusion(nn.Module):
    def __init__(self, cnn_dim=1536, vit_dim=768, fusion_dim=512, dropout=0.3):
        super().__init__()
        self.cnn_proj = nn.Linear(cnn_dim, fusion_dim)
        self.vit_proj = nn.Linear(vit_dim, fusion_dim)
        self.attn = nn.Sequential(
            nn.Linear(fusion_dim * 2, fusion_dim),
            nn.ReLU(),
            nn.Linear(fusion_dim, 2),
            nn.Softmax(dim=-1),
        )
        self.norm = nn.LayerNorm(fusion_dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, cnn_feat, vit_feat):
        c = self.cnn_proj(cnn_feat)
        v = self.vit_proj(vit_feat)
        w = self.attn(torch.cat([c, v], dim=-1))
        fused = w[:, 0:1] * c + w[:, 1:2] * v
        fused = self.norm(fused)
        fused = self.drop(fused)
        return fused


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  MODEL UTAMA
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class BrainHybridModel(nn.Module):
    """
    Nama class TETAP 'BrainHybridModel' (biar main.py/explainability.py
    tidak perlu diubah), tapi ISI-nya sekarang identik dengan
    HybridViTEfficientNet di notebook training.
    """

    def __init__(self, num_classes: int = NUM_CLASSES,
                 efficientnet_variant: str = "b3",
                 vit_embed_dim: int = 768,
                 vit_num_heads: int = 12,
                 vit_num_layers: int = 6,
                 fusion_dim: int = 512,
                 dropout: float = 0.3,
                 freeze_backbone: bool = True):
        super().__init__()

        # 1. CNN backbone (EfficientNet-B3)
        if HAS_WEIGHTS:
            backbone = efficientnet_b3(weights=EfficientNet_B3_Weights.DEFAULT)
        else:
            backbone = efficientnet_b3(pretrained=True)
        self.features = backbone.features
        self.cnn_out = 1536

        # 2. ViT branch (custom, dibangun dari feature map CNN)
        self.patch_embed = PatchEmbedding(self.cnn_out, patch_size=1,
                                           embed_dim=vit_embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, vit_embed_dim))
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        num_patches = (IMG_SIZE // 32) ** 2
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, vit_embed_dim))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        self.pos_drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(vit_embed_dim, vit_num_heads, dropout=dropout)
            for _ in range(vit_num_layers)
        ])
        self.vit_norm = nn.LayerNorm(vit_embed_dim)

        # 3. Fusion
        self.fusion = CrossModalAttentionFusion(
            cnn_dim=self.cnn_out, vit_dim=vit_embed_dim,
            fusion_dim=fusion_dim, dropout=dropout)

        # 4. Classifier head
        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim, 256),
            nn.GELU(),
            nn.BatchNorm1d(256),
            nn.Dropout(dropout),
            nn.Linear(256, num_classes),
        )

        # Freeze/unfreeze backbone (opsional, hanya relevan kalau fine-tune ulang)
        if freeze_backbone:
            for param in self.features.parameters():
                param.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        feat_map = self.features(x)
        cnn_feat = F.adaptive_avg_pool2d(feat_map, 1).flatten(1)
        patches = self.patch_embed(feat_map)
        cls = self.cls_token.expand(x.size(0), -1, -1)
        tokens = torch.cat([cls, patches], dim=1)
        tokens = tokens + self.pos_embed
        tokens = self.pos_drop(tokens)
        for blk in self.blocks:
            tokens = blk(tokens)
        tokens = self.vit_norm(tokens)
        vit_feat = tokens[:, 0]
        fused = self.fusion(cnn_feat, vit_feat)
        logits = self.classifier(fused)
        return logits

    def forward_with_attention(self, x: torch.Tensor):
        """
        Sama seperti forward(), tapi juga mengembalikan attention weights
        dari layer Transformer TERAKHIR -- dipakai untuk bikin heatmap
        'Peta Atensi Model' di explainability.py.

        Return:
            logits: [B, num_classes]
            last_attn: [B, num_heads, seq_len, seq_len]
        """
        feat_map = self.features(x)
        cnn_feat = F.adaptive_avg_pool2d(feat_map, 1).flatten(1)
        patches = self.patch_embed(feat_map)
        cls = self.cls_token.expand(x.size(0), -1, -1)
        tokens = torch.cat([cls, patches], dim=1)
        tokens = tokens + self.pos_embed
        tokens = self.pos_drop(tokens)

        last_attn = None
        for i, blk in enumerate(self.blocks):
            if i == len(self.blocks) - 1:
                tokens, last_attn = blk(tokens, return_attn=True)
            else:
                tokens = blk(tokens)
        tokens = self.vit_norm(tokens)
        vit_feat = tokens[:, 0]
        fused = self.fusion(cnn_feat, vit_feat)
        logits = self.classifier(fused)
        return logits, last_attn


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  UTILITIES
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def count_parameters(model: nn.Module):
    total     = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable, total - trainable


def print_model_info(model: nn.Module, device: str):
    try:
        from src.config import BATCH_SIZE, LR, EPOCHS, IMG_SIZE
    except ImportError:
        BATCH_SIZE, LR, EPOCHS, IMG_SIZE = 16, 5e-5, 30, 224

    total, trainable, frozen = count_parameters(model)
    sep70 = "=" * 70
    logger.info("\n" + sep70)
    logger.info("  INFORMASI MODEL")
    logger.info(sep70)
    col = 30
    fields = [
        ("Model Name",              "EfficientNet-B3 + Custom Vision Transformer (Hybrid)"),
        ("Architecture",            "EfficientNet-B3 Features + Custom ViT (6 layer) -> Cross-Modal Fusion -> MLP Head"),
        ("Jumlah Parameter",        f"{total:,}"),
        ("Trainable Parameter",     f"{trainable:,}"),
        ("Non-Trainable Parameter", f"{frozen:,}"),
        ("Image Size",              f"{IMG_SIZE} x {IMG_SIZE} px"),
        ("Batch Size",              str(BATCH_SIZE)),
        ("Learning Rate",           str(LR)),
        ("Epoch",                   str(EPOCHS)),
        ("Device",                  device.upper()),
    ]
    for k, v in fields:
        logger.info(f"  {k:<{col}}: {v}")
    logger.info("")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  QUICK SANITY CHECK
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

if __name__ == "__main__":
    model = BrainHybridModel()
    dummy = torch.randn(2, 3, IMG_SIZE, IMG_SIZE)
    out = model(dummy)
    print(f"Hybrid Model OK! Output shape: {out.shape}")
    total, trainable, frozen = count_parameters(model)
    print(f"Total: {total:,} | Trainable: {trainable:,} | Frozen: {frozen:,}")