src/analytics/xai.py

"""
SHAP 기반 XAI (설명 가능한 AI) 모듈.

XGBoost 분류 모델을 훈련하여 학교 상태(정상/주의/위험/고위험)를 분류하고,
SHAP TreeExplainer 로 각 피처의 기여도를 산출합니다.

피처 목록:
  - student_count_trend   : 최근 N년 학생 수 추세 (선형 기울기)
  - student_count_latest  : 최신 학생 수
  - teacher_student_ratio : 교사 1인당 학생 수
  - temp_teacher_ratio    : 기간제 교원 비율
  - transfer_net_avg      : 전출입 순 변화 평균
  - facility_age          : 시설 노후도 (현재 연도 - 개교 연도)
  - population_risk_index : 지역 소멸 위험 지수 (여성 20~39세 / 전체 노인 인구)
  - data_quality_score    : 데이터 신뢰도 점수
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# 상태 레이블
STATUS_LABELS = {0: "정상", 1: "주의", 2: "위험", 3: "고위험"}
STATUS_THRESHOLDS = {
    "student_count_latest": [60, 30, 10],   # > 60 → 정상, 30~60 → 주의, ...
}

FEATURE_COLS = [
    "student_count_trend",
    "student_count_latest",
    "teacher_student_ratio",
    "temp_teacher_ratio",
    "transfer_net_avg",
    "facility_age",
    "population_risk_index",
    "data_quality_score",
]

FEATURE_KR_NAMES = {
    "student_count_trend": "학생 수 추세",
    "student_count_latest": "현재 학생 수",
    "teacher_student_ratio": "교사 1인당 학생 수",
    "temp_teacher_ratio": "기간제 교원 비율",
    "transfer_net_avg": "전출입 순 변화",
    "facility_age": "시설 노후도",
    "population_risk_index": "지역 소멸 위험 지수",
    "data_quality_score": "데이터 신뢰도",
}


@dataclass
class ShapResult:
    """SHAP 분석 결과 컨테이너."""

    schul_code: str
    status_label: str
    status_code: int
    shap_values: dict[str, float]
    top_factors: list[dict[str, Any]]
    base_value: float = 0.0
    model_accuracy: float | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "schul_code": self.schul_code,
            "status_label": self.status_label,
            "status_code": self.status_code,
            "shap_values": self.shap_values,
            "top_factors": self.top_factors,
            "base_value": self.base_value,
            "model_accuracy": self.model_accuracy,
        }


def _compute_features(school_series: dict[str, Any]) -> dict[str, float]:
    """
    단일 학교 데이터 딕셔너리에서 SHAP 피처를 계산합니다.

    Parameters
    ----------
    school_series:
        키는 피처 이름, 값은 해당 지표. 연도별 시계열은
        "student_count_history" (list[float]) 형태로 전달합니다.

    Returns
    -------
    dict[str, float]
        FEATURE_COLS 에 해당하는 피처 값 딕셔너리.
    """
    features: dict[str, float] = {}

    # 학생 수 히스토리에서 추세(기울기) 산출
    history = school_series.get("student_count_history", [])
    if len(history) >= 2:
        x = np.arange(len(history), dtype=float)
        slope = float(np.polyfit(x, history, 1)[0])
        features["student_count_trend"] = slope
        features["student_count_latest"] = float(history[-1])
    elif len(history) == 1:
        features["student_count_trend"] = 0.0
        features["student_count_latest"] = float(history[-1])
    else:
        features["student_count_trend"] = 0.0
        features["student_count_latest"] = float(school_series.get("student_count", 0) or 0)

    # 교사 1인당 학생 수
    teacher_count = float(school_series.get("teacher_count", 1) or 1)
    student_count = features["student_count_latest"]
    features["teacher_student_ratio"] = student_count / max(teacher_count, 1)

    # 기간제 교원 비율
    temp_teacher = float(school_series.get("temp_teacher_count", 0) or 0)
    features["temp_teacher_ratio"] = temp_teacher / max(teacher_count, 1)

    # 전출입 순 변화 평균
    features["transfer_net_avg"] = float(school_series.get("transfer_net_avg", 0) or 0)

    # 시설 노후도
    established = int(school_series.get("established_year", 0) or 0)
    current_year = pd.Timestamp.now().year
    features["facility_age"] = float(current_year - established) if established > 0 else float("nan")

    # 지역 소멸 위험 지수 (여성 20~39세 / (65세 이상 인구 * 0.5))
    women_20_39 = float(school_series.get("region_women_20_39", 0) or 0)
    elderly = float(school_series.get("region_elderly_pop", 1) or 1)
    features["population_risk_index"] = women_20_39 / max(elderly * 0.5, 1)

    # 데이터 신뢰도
    features["data_quality_score"] = float(school_series.get("data_quality_score", 1.0) or 1.0)

    return features


def _rule_based_status(features: dict[str, float]) -> int:
    """
    규칙 기반 상태 분류 (학습 데이터 부족 시 fallback).

    Returns
    -------
    int
        0=정상, 1=주의, 2=위험, 3=고위험
    """
    student = features.get("student_count_latest", 100)
    trend = features.get("student_count_trend", 0)
    risk = features.get("population_risk_index", 1.0)

    score = 0
    if student < 10:
        score += 3
    elif student < 30:
        score += 2
    elif student < 60:
        score += 1

    if trend < -10:
        score += 2
    elif trend < -5:
        score += 1

    if risk < 0.5:
        score += 1

    return min(score, 3)


class ShapExplainer:
    """
    XGBoost + SHAP TreeExplainer 기반 학교 상태 분류 및 설명 엔진.

    - 훈련 데이터가 충분하면 XGBoost 를 학습하고 SHAP 값을 산출합니다.
    - 훈련 데이터가 부족하면 규칙 기반 분류 + KernelExplainer fallback 을 사용합니다.

    사용 예::

        explainer = ShapExplainer()
        explainer.fit(train_df)

        result = explainer.explain(
            schul_code="7431234",
            school_data={
                "student_count_history": [120, 105, 98, 87, 75],
                "teacher_count": 6,
                "temp_teacher_count": 1,
                "established_year": 1975,
                ...
            }
        )
    """

    _MIN_TRAIN_SAMPLES = 20  # 모델 학습에 필요한 최소 샘플 수

    def __init__(self) -> None:
        self._model: Any = None
        self._explainer: Any = None
        self._is_fitted = False
        self._model_accuracy: float | None = None

    def fit(self, train_df: pd.DataFrame, label_col: str = "status_code") -> "ShapExplainer":
        """
        XGBoost 분류 모델을 훈련합니다.

        Parameters
        ----------
        train_df:
            FEATURE_COLS + label_col 을 포함하는 DataFrame.
            label_col 은 0~3 정수 레이블.
        label_col:
            분류 대상 레이블 컬럼명.

        Returns
        -------
        self
        """
        available_features = [c for c in FEATURE_COLS if c in train_df.columns]
        if len(train_df) < self._MIN_TRAIN_SAMPLES or not available_features:
            logger.warning(
                "훈련 샘플 부족 (%d개). 규칙 기반 분류를 사용합니다.", len(train_df)
            )
            return self

        try:
            import shap  # type: ignore[import]
            from sklearn.model_selection import train_test_split  # type: ignore[import]
            from xgboost import XGBClassifier  # type: ignore[import]

            X = train_df[available_features].fillna(0)
            y = train_df[label_col]

            X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

            model = XGBClassifier(
                n_estimators=100,
                max_depth=4,
                learning_rate=0.1,
                use_label_encoder=False,
                eval_metric="mlogloss",
                random_state=42,
            )
            model.fit(X_train, y_train)
            self._model_accuracy = float(model.score(X_val, y_val))

            self._model = model
            self._explainer = shap.TreeExplainer(model)
            self._is_fitted = True

            logger.info(
                "XGBoost 학습 완료: samples=%d accuracy=%.3f",
                len(train_df),
                self._model_accuracy,
            )
        except Exception as exc:  # noqa: BLE001
            logger.error("XGBoost / SHAP 초기화 실패: %s", exc)

        return self

    def explain(
        self,
        schul_code: str,
        school_data: dict[str, Any],
    ) -> ShapResult:
        """
        단일 학교에 대한 상태 분류 및 SHAP 기여도를 산출합니다.

        Parameters
        ----------
        schul_code:
            대상 학교 SD_SCHUL_CODE.
        school_data:
            _compute_features 가 이해하는 학교 지표 딕셔너리.

        Returns
        -------
        ShapResult
        """
        features = _compute_features(school_data)
        feature_vector = pd.DataFrame([features])[FEATURE_COLS].fillna(0)

        if self._is_fitted and self._model is not None and self._explainer is not None:
            status_code = int(self._model.predict(feature_vector)[0])
            try:
                shap_vals = self._explainer.shap_values(feature_vector)
                # multi-class → 예측 클래스의 SHAP 배열 선택
                if isinstance(shap_vals, list):
                    sv = shap_vals[status_code][0]
                else:
                    sv = shap_vals[0]
                base_val = float(
                    self._explainer.expected_value[status_code]
                    if isinstance(self._explainer.expected_value, (list, np.ndarray))
                    else self._explainer.expected_value
                )
            except Exception as exc:  # noqa: BLE001
                logger.warning("SHAP 값 산출 실패, 0 벡터 사용: %s", exc)
                sv = np.zeros(len(FEATURE_COLS))
                base_val = 0.0
        else:
            # fallback: 규칙 기반 분류
            status_code = _rule_based_status(features)
            sv = np.zeros(len(FEATURE_COLS))
            base_val = 0.0

        shap_dict = {col: round(float(v), 4) for col, v in zip(FEATURE_COLS, sv)}

        # 상위 기여 요인 (절댓값 기준 정렬)
        sorted_factors = sorted(shap_dict.items(), key=lambda x: abs(x[1]), reverse=True)
        top_factors = [
            {
                "feature": col,
                "feature_kr": FEATURE_KR_NAMES.get(col, col),
                "shap_value": val,
                "feature_value": round(float(features.get(col, 0)), 4),
            }
            for col, val in sorted_factors[:5]
        ]

        logger.info(
            "SHAP 분석 완료: schul_code=%s status=%s top_factor=%s",
            schul_code,
            STATUS_LABELS[status_code],
            sorted_factors[0][0] if sorted_factors else "N/A",
        )

        return ShapResult(
            schul_code=schul_code,
            status_label=STATUS_LABELS[status_code],
            status_code=status_code,
            shap_values=shap_dict,
            top_factors=top_factors,
            base_value=base_val,
            model_accuracy=self._model_accuracy,
        )