"""
model.py
========
Định nghĩa kiến trúc mô hình ResNet-50 cho 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 ResNet50_DR(nn.Module):
    """
    Kiến trúc mô hình ResNet-50 cho phân loại 5 lớp DR:
    - Backbone: ResNet-50 (Feature Extractor: 2048 channels)
    - Classifier: Dropout(drop_rate) + Linear(2048 -> 5 classes)
    """
    def __init__(
        self,
        num_classes: int = 5,
        drop_rate: float = 0.3,
        pretrained: bool = False,
    ):
        super().__init__()
        if pretrained:
            weights = models.ResNet50_Weights.DEFAULT
            self.model = models.resnet50(weights=weights)
        else:
            self.model = models.resnet50(weights=None)

        in_features = self.model.fc.in_features  # 2048 channels
        self.model.fc = nn.Sequential(
            nn.Dropout(p=drop_rate),
            nn.Linear(in_features, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def load_state_dict(self, state_dict: dict, strict: bool = True):
        """Hỗ trợ nạp cả state_dict của torchvision.resnet50 chuẩn lẫn wrapped module."""
        if any(k.startswith("model.") for k in state_dict.keys()):
            return super().load_state_dict(state_dict, strict=strict)
        else:
            return self.model.load_state_dict(state_dict, strict=strict)


def build_model(
    num_classes: int = 5,
    drop_rate: float = 0.3,
    pretrained: bool = False,
) -> ResNet50_DR:
    """Hàm helper khởi tạo mô hình ResNet50_DR."""
    model = ResNet50_DR(
        num_classes=num_classes,
        drop_rate=drop_rate,
        pretrained=pretrained,
    )
    return model