"""
P3 (Fashion Image Segmentation · Attribute Tagging) -- 공통 유틸리티

2단계(2-Stage) 파이프라인 전반에서 공유하는 보조 도구를 모았습니다.
    Stage 1 : Mask R-CNN instance segmentation (46개 카테고리)
    Stage 2 : 인스턴스 crop 기반 multi-label attribute tagging (294개 속성)

[주요 구성]
    1. 카테고리 상수      MAIN_GARMENT_NAMES(27) / GARMENT_PART_NAMES(19)
    2. load_category_ids()        런타임에 COCO JSON → (MAIN_IDS, PART_IDS, id2name)
    3. load_attribute_supers()    인덱스 → 속성 슈퍼카테고리 매핑 (그룹별 메트릭용)
    4. set_seed()                 재현성 시드 고정
    5. AverageMeter               실행 평균 추적
    6. EarlyStopping              조기 종료 (mAP@0.5 / F1 모니터링)
    7. compute_segmentation_metrics()   Stage 1 -- COCO mAP (Main/Part 분리)
    8. compute_attribute_metrics()      Stage 2 -- macro F1 + 슈퍼카테고리별 F1
    9. collate_fn_stage1()        Mask R-CNN용 variable-size collate
   10. setup_logger()            파일 + 콘솔 동시 로깅
   11. get_pos_weight()         inverse-frequency BCE pos_weight

근거: 카테고리/속성 구조는 EDA(`004_EDA/P3_EDA_short.ipynb` cell-05)에서 검증.
"""

from __future__ import annotations

import os
import json
import random
import logging
from collections import Counter
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch


# ============================================================================
# 1. 카테고리 상수 (EDA 결과 기반 -- 004_EDA/P3_EDA_short.ipynb cell-05)
# ----------------------------------------------------------------------------
# Fashionpedia 46개 카테고리 = Main 의류 27개 + Part 부속/장식 19개.
# Main 만 Stage 2 속성 태깅의 입력으로 사용한다 (EDA 가정 ③: Main 속성이
# Part 보다 압도적으로 풍부). 이름은 COCO `categories[*].name` 과 정확히 일치.
# ============================================================================

#: Stage 2 입력이 되는 본체 의류 27개 (옷 + 액세서리)
MAIN_GARMENT_NAMES = {
    # 상의
    "shirt, blouse", "top, t-shirt, sweatshirt", "sweater", "cardigan",
    "jacket", "vest",
    # 하의 / 원피스류
    "pants", "shorts", "skirt", "coat", "dress", "jumpsuit", "cape",
    # 헤드 / 액세서리
    "glasses", "hat", "headband, head covering, hair accessory",
    "tie", "glove", "watch", "belt",
    # 다리 / 발
    "leg warmer", "tights, stockings", "sock", "shoe",
    # 소품
    "bag, wallet", "scarf", "umbrella",
}

#: 의류에 부속/장식되는 파트 19개 (Stage 2에서 제외, Stage 1에서만 검출)
GARMENT_PART_NAMES = {
    # 구조 부위
    "hood", "collar", "lapel", "epaulette", "sleeve", "pocket", "neckline",
    "buckle", "zipper",
    # 장식 요소
    "applique", "bead", "bow", "flower", "fringe", "ribbon", "rivet",
    "ruffle", "sequin", "tassel",
}


def load_category_ids(
    coco_json_path: str,
) -> Tuple[set, set, Dict[int, str]]:
    """COCO JSON을 읽어 Main / Part 카테고리 ID 집합과 ID→이름 매핑을 만든다.

    카테고리 ID 는 데이터셋 버전에 따라 달라질 수 있으므로 하드코딩하지 않고
    런타임에 이름(MAIN_GARMENT_NAMES / GARMENT_PART_NAMES)으로 매핑한다.

    Args:
        coco_json_path (str): instances_attributes_*.json 경로

    Returns:
        Tuple[set, set, Dict[int, str]]:
            - MAIN_IDS : Main 카테고리 ID 집합 (27개)
            - PART_IDS : Part 카테고리 ID 집합 (19개)
            - id2name  : {category_id: name} 전체 매핑 (46개)

    Raises:
        ValueError: JSON 의 카테고리 이름이 정의된 상수와 어긋날 때
            (Main+Part 합이 46이 아니거나 분류되지 않은 이름이 있을 때).
    """
    with open(coco_json_path, "r") as f:
        raw = json.load(f)

    id2name = {c["id"]: c["name"] for c in raw["categories"]}

    main_ids, part_ids, unknown = set(), set(), []
    for cid, name in id2name.items():
        if name in MAIN_GARMENT_NAMES:
            main_ids.add(cid)
        elif name in GARMENT_PART_NAMES:
            part_ids.add(cid)
        else:
            unknown.append(name)

    if unknown:
        raise ValueError(
            f"Main/Part 어느 쪽에도 속하지 않는 카테고리: {unknown}. "
            "MAIN_GARMENT_NAMES / GARMENT_PART_NAMES 상수를 확인하세요."
        )
    if len(main_ids) != len(MAIN_GARMENT_NAMES) or len(part_ids) != len(GARMENT_PART_NAMES):
        raise ValueError(
            f"카테고리 개수 불일치: Main {len(main_ids)}/{len(MAIN_GARMENT_NAMES)}, "
            f"Part {len(part_ids)}/{len(GARMENT_PART_NAMES)}"
        )

    return main_ids, part_ids, id2name


def load_attribute_supers(
    coco_json_path: str,
    num_attrs: int = 294,
) -> Tuple[List[str], Dict[int, int]]:
    """속성 인덱스 → 슈퍼카테고리 매핑을 만든다 (슈퍼카테고리별 F1 메트릭용).

    Fashionpedia 속성 ID 는 0~340 범위에 비연속적으로 분포(총 294개)하므로,
    학습용 multi-hot 벡터의 인덱스는 **속성 ID 오름차순 순서**로 0..293에 매핑한다.
    preprocess 의 multi-hot 생성과 반드시 동일한 정렬 규칙을 사용해야 한다.

    Args:
        coco_json_path (str): instances_attributes_*.json 경로
        num_attrs (int): 속성 개수 (기본 294)

    Returns:
        Tuple[List[str], Dict[int, int]]:
            - idx2super : 길이 num_attrs 리스트. idx2super[i] = i번째 속성의 슈퍼카테고리
            - attr_id2idx : {원본 attribute_id: 0..num_attrs-1 인덱스}

    Note:
        Fashionpedia 데이터의 실제 슈퍼카테고리는 11종이다
        (animal, leather, length, neckline type, nickname, non-textile material type,
         opening type, silhouette, textile finishing, textile pattern, waistline).
    """
    with open(coco_json_path, "r") as f:
        raw = json.load(f)

    attrs = sorted(raw["attributes"], key=lambda a: a["id"])
    if len(attrs) != num_attrs:
        raise ValueError(f"속성 개수 불일치: JSON {len(attrs)} != num_attrs {num_attrs}")

    idx2super = [a["supercategory"] for a in attrs]
    attr_id2idx = {a["id"]: i for i, a in enumerate(attrs)}
    return idx2super, attr_id2idx


# ============================================================================
# 2. 재현성
# ============================================================================

def set_seed(seed: int = 42) -> None:
    """random / numpy / torch / cudnn 시드를 모두 고정해 재현성을 확보한다.

    Args:
        seed (int): 랜덤 시드 (기본 42)
    """
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # 완전 결정론적 동작 (속도보다 재현성 우선)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# ============================================================================
# 3. 학습 보조 클래스
# ============================================================================

class AverageMeter:
    """실행 평균과 합계를 추적하는 누적기.

    배치 단위 loss / 메트릭의 에폭 평균을 구할 때 사용한다.

    사용 예:
        meter = AverageMeter("loss")
        for loss, bs in batches:
            meter.update(loss, n=bs)
        print(meter.avg)
    """

    def __init__(self, name: str = "metric") -> None:
        self.name = name
        self.reset()

    def reset(self) -> None:
        """모든 통계를 0으로 초기화."""
        self.val = 0.0
        self.avg = 0.0
        self.sum = 0.0
        self.count = 0

    def update(self, val: float, n: int = 1) -> None:
        """새 값으로 통계를 갱신한다.

        Args:
            val (float): 현재 배치의 (평균) 값
            n (int): 표본 수 (배치 크기). 가중 평균에 사용
        """
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count if self.count > 0 else 0.0

    def __repr__(self) -> str:
        return f"{self.name}: {self.avg:.4f} (n={self.count})"


class EarlyStopping:
    """검증 메트릭이 patience 에폭 동안 개선되지 않으면 학습 중단 신호를 준다.

    Stage 1 은 mAP@0.5 (mode='max'), Stage 2 는 macro F1 (mode='max') 을 모니터링.
    loss 처럼 작아야 좋은 지표는 mode='min' 으로 사용.
    """

    def __init__(
        self,
        patience: int = 5,
        mode: str = "max",
        delta: float = 0.0,
    ) -> None:
        """초기화.

        Args:
            patience (int): 개선 없이 허용할 최대 에폭 수
            mode (str): "max"(클수록 좋음) 또는 "min"(작을수록 좋음)
            delta (float): 개선으로 인정할 최소 변화량
        """
        if mode not in ("max", "min"):
            raise ValueError(f"mode 는 'max' 또는 'min' 이어야 합니다 (받은 값: {mode})")
        self.patience = patience
        self.mode = mode
        self.delta = delta
        self.counter = 0
        self.best_score: Optional[float] = None
        self.early_stop = False

    def _is_improved(self, score: float) -> bool:
        if self.best_score is None:
            return True
        if self.mode == "max":
            return score > self.best_score + self.delta
        return score < self.best_score - self.delta

    def __call__(self, metric: float) -> bool:
        """현재 에폭 메트릭을 받아 학습 중단 여부를 반환한다.

        Args:
            metric (float): 현재 에폭의 검증 메트릭

        Returns:
            bool: True 이면 학습을 멈춰야 함
        """
        if self._is_improved(metric):
            self.best_score = metric
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        return self.early_stop


# ============================================================================
# 4. collate_fn (Stage 1)
# ============================================================================

def collate_fn_stage1(batch: List[Tuple]) -> Tuple[List, List]:
    """Mask R-CNN용 collate 함수.

    torchvision detection 모델은 서로 다른 크기의 이미지를 하나의 텐서로
    stack 하지 않고 **리스트** 그대로 입력받는다. 기본 collate 의 자동 stacking
    을 우회하기 위해 (이미지 리스트, 타겟 리스트) 로만 분리한다.

    Args:
        batch (List[Tuple]): [(image_tensor, target_dict), ...]

    Returns:
        Tuple[List, List]: (images, targets) -- 둘 다 길이 B 리스트
    """
    images, targets = zip(*batch)
    return list(images), list(targets)


# ============================================================================
# 5. Stage 1 메트릭 -- COCO mAP
# ============================================================================

def compute_segmentation_metrics(
    coco_gt,
    coco_dt,
    main_ids: Optional[set] = None,
    part_ids: Optional[set] = None,
    iou_type: str = "segm",
) -> Dict[str, float]:
    """pycocotools COCOeval 로 instance segmentation mAP 를 계산한다.

    Args:
        coco_gt: GT COCO 객체 (`pycocotools.coco.COCO`)
        coco_dt: 예측 결과 COCO 객체. 보통 `coco_gt.loadRes(results)` 로 생성하며,
            results 는 [{"image_id", "category_id", "segmentation"(RLE),
            "score"}, ...] 형식의 리스트.
        main_ids (set, optional): Main 카테고리 ID 집합. 주면 Main mAP 분리 계산.
        part_ids (set, optional): Part 카테고리 ID 집합. 주면 Part mAP 분리 계산.
        iou_type (str): "segm"(마스크) 또는 "bbox".

    Returns:
        Dict[str, float]: {
            "mAP_50":      mAP @ IoU=0.50,
            "mAP_50_95":   mAP @ IoU=0.50:0.95 (COCO primary),
            "main_mAP_50": (main_ids 제공 시) Main 카테고리 mAP@0.5,
            "part_mAP_50": (part_ids 제공 시) Part 카테고리 mAP@0.5,
        }

    Note:
        COCOeval.stats[0] = mAP@[.5:.95], stats[1] = mAP@0.5.
        값이 -1 이면 해당 카테고리에 대한 예측/GT 가 없다는 뜻.
    """
    from pycocotools.cocoeval import COCOeval

    def _run(cat_ids: Optional[Sequence[int]] = None) -> Tuple[float, float]:
        evaluator = COCOeval(coco_gt, coco_dt, iouType=iou_type)
        if cat_ids is not None:
            evaluator.params.catIds = list(cat_ids)
        evaluator.evaluate()
        evaluator.accumulate()
        evaluator.summarize()
        # stats[0]=mAP@[.5:.95], stats[1]=mAP@0.5
        return float(evaluator.stats[1]), float(evaluator.stats[0])

    map_50, map_50_95 = _run()
    metrics: Dict[str, float] = {"mAP_50": map_50, "mAP_50_95": map_50_95}

    if main_ids:
        metrics["main_mAP_50"] = _run(main_ids)[0]
    if part_ids:
        metrics["part_mAP_50"] = _run(part_ids)[0]

    return metrics


# ============================================================================
# 6. Stage 2 메트릭 -- multi-label attribute F1
# ============================================================================

def compute_attribute_metrics(
    logits,
    targets,
    threshold: float = 0.5,
    attr_supers: Optional[Sequence[str]] = None,
) -> Dict[str, float]:
    """multi-label 속성 분류 메트릭 (macro F1/precision/recall + 슈퍼카테고리별 F1).

    Args:
        logits: 모델 출력 [N, num_attrs]. **raw logit** 으로 보고 내부에서
            sigmoid 를 적용한다 (torch.Tensor / np.ndarray 모두 허용).
        targets: multi-hot 정답 [N, num_attrs] (0/1).
        threshold (float): sigmoid 확률을 1로 판정할 임계값.
        attr_supers (Sequence[str], optional): 길이 num_attrs. 각 속성 인덱스의
            슈퍼카테고리 이름 (load_attribute_supers 로 얻음). 주면 그룹별
            macro F1 을 함께 반환한다.

    Returns:
        Dict[str, float]: {
            "f1_macro", "precision_macro", "recall_macro",
            "f1_micro",
            "f1_super/<supercategory>": ...  (attr_supers 제공 시 그룹별)
        }
    """
    from sklearn.metrics import f1_score, precision_score, recall_score

    probs = _to_numpy(logits)
    y_true = _to_numpy(targets).astype(int)
    y_pred = (_sigmoid(probs) >= threshold).astype(int)

    metrics = {
        "f1_macro": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
        "precision_macro": float(precision_score(y_true, y_pred, average="macro", zero_division=0)),
        "recall_macro": float(recall_score(y_true, y_pred, average="macro", zero_division=0)),
        "f1_micro": float(f1_score(y_true, y_pred, average="micro", zero_division=0)),
    }

    # 슈퍼카테고리(9~11개 그룹)별 macro F1
    if attr_supers is not None:
        if len(attr_supers) != y_true.shape[1]:
            raise ValueError(
                f"attr_supers 길이({len(attr_supers)})가 속성 차원({y_true.shape[1]})과 다릅니다."
            )
        attr_supers = np.asarray(attr_supers)
        for group in sorted(set(attr_supers.tolist())):
            cols = np.where(attr_supers == group)[0]
            f1_g = f1_score(
                y_true[:, cols], y_pred[:, cols], average="macro", zero_division=0
            )
            metrics[f"f1_super/{group}"] = float(f1_g)

    return metrics


# ============================================================================
# 7. 로깅
# ============================================================================

def setup_logger(log_dir: str, name: str = "train") -> logging.Logger:
    """파일과 콘솔에 동시 출력하는 로거를 만든다.

    Args:
        log_dir (str): 로그 파일을 저장할 디렉토리 (없으면 생성)
        name (str): 로거 이름 겸 로그 파일명 (`<name>.log`)

    Returns:
        logging.Logger: 설정된 로거
    """
    os.makedirs(log_dir, exist_ok=True)
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)

    # 중복 핸들러 방지 (재호출 시 로그가 여러 번 찍히는 문제)
    if logger.handlers:
        return logger
    logger.propagate = False

    fmt = logging.Formatter(
        "[%(asctime)s] %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
    )

    file_handler = logging.FileHandler(os.path.join(log_dir, f"{name}.log"))
    file_handler.setFormatter(fmt)
    logger.addHandler(file_handler)

    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(fmt)
    logger.addHandler(stream_handler)

    return logger


# ============================================================================
# 8. BCE pos_weight (속성 long-tail 보정 -- EDA 가정 ⑤)
# ============================================================================

def get_pos_weight(
    attr_counter: Counter,
    num_attrs: int = 294,
    max_weight: float = 10.0,
    attr_ids: Optional[Sequence[int]] = None,
) -> torch.Tensor:
    """속성 빈도의 역수로 BCEWithLogitsLoss 의 pos_weight 를 계산한다.

    희귀 속성(long-tail)일수록 양성 신호가 약하므로 손실에서 더 크게 가중한다.
    가장 흔한 속성을 기준(1.0)으로 한 inverse-frequency 가중을 사용하고,
    극단적 가중을 막기 위해 max_weight 로 clipping 한다.

        pos_weight[i] = clip(max(count) / count[i],  1.0,  max_weight)

    Args:
        attr_counter (Counter): {속성키: 등장횟수}. 키는 attr_ids 와 같은 체계여야 함.
        num_attrs (int): 속성 개수 (출력 벡터 길이, 기본 294)
        max_weight (float): pos_weight 상한 (기본 10.0)
        attr_ids (Sequence[int], optional): 인덱스 i 에 대응하는 원본 속성 ID 목록.
            Fashionpedia 처럼 ID 가 비연속(0~340)일 때 정렬된 ID 리스트를 넘긴다.
            None 이면 키를 0..num_attrs-1 인덱스로 간주(연속 가정).

    Returns:
        torch.Tensor: shape (num_attrs,), dtype float32.
    """
    keys = list(attr_ids) if attr_ids is not None else list(range(num_attrs))
    if len(keys) != num_attrs:
        raise ValueError(f"attr_ids 길이({len(keys)})가 num_attrs({num_attrs})와 다릅니다.")

    # 등장 0회는 1로 바닥 처리 (0 나눗셈 방지)
    counts = np.array([max(attr_counter.get(k, 0), 1) for k in keys], dtype=np.float64)
    pos_weight = counts.max() / counts
    pos_weight = np.clip(pos_weight, 1.0, max_weight)

    return torch.tensor(pos_weight, dtype=torch.float32)


# ============================================================================
# 내부 헬퍼
# ============================================================================

def _to_numpy(x) -> np.ndarray:
    """torch.Tensor / np.ndarray / list 를 numpy 배열로 변환."""
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    return np.asarray(x)


def _sigmoid(x: np.ndarray) -> np.ndarray:
    """수치 안정적인 sigmoid."""
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))