"""
model.py
========
Định nghĩa kiến trúc mô hình EfficientNet-B4 + CBAM (Convolutional Block Attention Module)
phục vụ bài toán phân loại Mức độ Bệnh Võng mạc Tiểu đường (Diabetic Retinopathy - 5 lớp ICDR).
"""

from __future__ import annotations
import torch
import torch.nn as nn
from torchvision import models


class ChannelAttention(nn.Module):
    """
    Channel Attention Sub-module của CBAM.
    Tính toán trọng số chú ý cho từng kênh đặc trưng dựa trên AvgPool và MaxPool.
    """
    def __init__(self, in_planes: int, ratio: int = 16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        return self.sigmoid(avg_out + max_out)


class SpatialAttention(nn.Module):
    """
    Spatial Attention Sub-module của CBAM.
    Tập trung vào các vùng không gian quan trọng (xuất huyết, vi phình mạch, xuất tiết).
    """
    def __init__(self, kernel_size: int = 7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))


class CBAM(nn.Module):
    """
    Convolutional Block Attention Module (CBAM).
    Kết hợp Channel Attention và Spatial Attention nối tiếp.
    """
    def __init__(self, in_planes: int, ratio: int = 16, kernel_size: int = 7):
        super().__init__()
        self.ca = ChannelAttention(in_planes, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x * self.ca(x)
        x = x * self.sa(x)
        return x


class EfficientNetB4_CBAM(nn.Module):
    """
    Kiến trúc mô hình chính:
    - Backbone: EfficientNet-B4 (Feature Extractor: 1792 channels)
    - Attention: CBAM Module đặt ngay sau feature extractor
    - Global Pooling: AdaptiveAvgPool2d(1)
    - Classifier: Dropout(0.3) + Linear(1792 -> 5 classes)
    """
    def __init__(
        self,
        num_classes: int = 5,
        drop_rate: float = 0.3,
        cbam_ratio: int = 16,
        pretrained: bool = False,
    ):
        super().__init__()
        if pretrained:
            weights = models.EfficientNet_B4_Weights.DEFAULT
            backbone = models.efficientnet_b4(weights=weights)
        else:
            backbone = models.efficientnet_b4(weights=None)

        self.features = backbone.features
        in_planes = 1792  # Output channel size of EfficientNet-B4 features
        self.cbam = CBAM(in_planes, ratio=cbam_ratio)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(p=drop_rate),
            nn.Linear(in_planes, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.cbam(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


def build_model(
    num_classes: int = 5,
    drop_rate: float = 0.3,
    cbam_ratio: int = 16,
    pretrained: bool = False,
) -> EfficientNetB4_CBAM:
    """Hàm helper khởi tạo mô hình và in số lượng tham số."""
    model = EfficientNetB4_CBAM(
        num_classes=num_classes,
        drop_rate=drop_rate,
        cbam_ratio=cbam_ratio,
        pretrained=pretrained,
    )
    return model