"""
P3 모델 -- Fashionpedia 2-Stage Model (Fashion Image Segmentation · Attribute Tagging)

[Stage 1] build_maskrcnn() -- Mask R-CNN ResNet50-FPN v2 (COCO pretrained)
    · box/mask predictor 를 Fashionpedia 47-class(46 카테고리 + background) 로 교체
    · 입력: List[Tensor[3,H,W]] / 출력: train→loss dict, eval→예측 List[Dict]

[Stage 2] AttributeClassifier -- ResNet-50 기반 multi-label 속성 분류
    · 입력: Tensor[B,3,224,224] (배경 제거된 crop)
    · 출력: logits [B,294]  (sigmoid 는 BCEWithLogitsLoss 내부에서)

[헬퍼] build_model(stage, **kwargs) -- 'maskrcnn' / 'attribute' 분기

[중요] Stage 1 label 은 dataset 에서 category_id+1 로 들어온다 (0=background).
       따라서 num_classes=47, 추론 결과를 보여줄 땐 label-1 로 되돌린다.
"""

from __future__ import annotations

from typing import List

import torch
import torch.nn as nn

import torchvision
from torchvision.models.detection import (
    maskrcnn_resnet50_fpn_v2,
    MaskRCNN_ResNet50_FPN_V2_Weights,
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor


# ============================================================================
# Stage 1: Mask R-CNN
# ============================================================================

def build_maskrcnn(num_classes: int = 47, pretrained: bool = True) -> nn.Module:
    """Mask R-CNN ResNet50-FPN v2 를 Fashionpedia 용으로 만든다.

    COCO 로 사전학습된 모델을 불러온 뒤, 마지막 분류/마스크 head 만 우리 클래스 수
    (47 = 46 카테고리 + background)에 맞게 갈아끼운다. backbone(특징 추출부)은
    COCO 가중치를 그대로 살려 전이학습한다.

    Args:
        num_classes (int): 클래스 수 (background 포함). Fashionpedia 는 47.
        pretrained (bool): COCO 사전학습 가중치 사용 여부.

    Returns:
        nn.Module: head 가 교체된 Mask R-CNN.
            · train(): model(images, targets) → loss dict
            · eval():  model(images)          → [{boxes,labels,scores,masks}, ...]
    """
    weights = MaskRCNN_ResNet50_FPN_V2_Weights.COCO_V1 if pretrained else None
    model = maskrcnn_resnet50_fpn_v2(weights=weights)

    # 1) 박스 분류 head 교체 (cls_score / bbox_pred 를 num_classes 로)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 2) 마스크 예측 head 교체
    in_channels_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_dim = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_channels_mask, hidden_dim, num_classes)

    return model


# ============================================================================
# Stage 2: Attribute Classifier
# ============================================================================

class AttributeClassifier(nn.Module):
    """ResNet-50 기반 multi-label 속성 분류기 (Stage 2).

    배경이 제거된 인스턴스 crop 을 받아 294개 속성에 대한 logit 을 출력한다.
    backbone(특징 추출부)은 ImageNet 사전학습 가중치를 쓰고, 그 위에 새 분류 head 를 얹는다.

    구조:
        backbone : ResNet-50 (GAP/fc 제거) → feature map [B, 2048, 7, 7]
        head     : GAP → Dropout → Linear(2048→512) → ReLU → Dropout → Linear(512→294)

    Args:
        backbone (str): 'resnet50' (그 외 이름은 timm 으로 생성 시도).
        num_attrs (int): 속성 수 (294).
        pretrained (bool): ImageNet 사전학습 사용 여부.
        dropout (float): 드롭아웃 비율.
    """

    def __init__(
        self,
        backbone: str = "resnet50",
        num_attrs: int = 294,
        pretrained: bool = True,
        dropout: float = 0.3,
    ) -> None:
        super().__init__()
        self.num_attrs = num_attrs

        if backbone == "resnet50":
            from torchvision.models import resnet50, ResNet50_Weights
            weights = ResNet50_Weights.IMAGENET1K_V2 if pretrained else None
            net = resnet50(weights=weights)
            self.feature_dim = net.fc.in_features                 # 2048
            # 마지막 avgpool(GAP) + fc 를 떼어내 feature map 추출부만 남긴다
            self.backbone = nn.Sequential(*list(net.children())[:-2])
        else:
            import timm
            # global_pool='' → GAP 미적용 feature map, num_classes=0 → fc 제거
            self.backbone = timm.create_model(
                backbone, pretrained=pretrained, num_classes=0, global_pool=""
            )
            self.feature_dim = self.backbone.num_features

        # 분류 head (GAP 포함). sigmoid 는 적용하지 않음 (BCEWithLogitsLoss).
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(1),
            nn.Dropout(dropout),
            nn.Linear(self.feature_dim, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(512, num_attrs),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """순전파.

        Args:
            x (torch.Tensor): [B, 3, 224, 224]

        Returns:
            torch.Tensor: logits [B, num_attrs] (sigmoid 전)
        """
        features = self.backbone(x)        # [B, 2048, 7, 7]
        logits = self.head(features)       # [B, num_attrs]
        return logits

    def get_param_groups(
        self, backbone_lr: float = 1e-5, head_lr: float = 1e-4
    ) -> List[dict]:
        """backbone 과 head 에 서로 다른 학습률을 주는 파라미터 그룹.

        사전학습된 backbone 은 살짝만(작은 lr), 새로 얹은 head 는 빠르게(큰 lr) 학습한다.

        Args:
            backbone_lr (float): backbone 학습률.
            head_lr (float): head 학습률.

        Returns:
            List[dict]: optimizer 에 그대로 넘길 파라미터 그룹.
        """
        return [
            {"params": self.backbone.parameters(), "lr": backbone_lr},
            {"params": self.head.parameters(), "lr": head_lr},
        ]


# ============================================================================
# 빌더 헬퍼
# ============================================================================

def build_model(stage: str, **kwargs) -> nn.Module:
    """stage 이름으로 알맞은 모델을 만든다.

    Args:
        stage (str): 'maskrcnn'(Stage 1) 또는 'attribute'(Stage 2).
        **kwargs: 각 빌더로 그대로 전달 (num_classes/pretrained 또는
            backbone/num_attrs/pretrained/dropout 등).

    Returns:
        nn.Module: 생성된 모델.
    """
    if stage == "maskrcnn":
        return build_maskrcnn(**kwargs)
    if stage == "attribute":
        return AttributeClassifier(**kwargs)
    raise ValueError(f"알 수 없는 stage: {stage!r} (가능: 'maskrcnn', 'attribute')")