"""DM 개인별 위험요인 - local SHAP.

LLM_Report/utils/dm_risk.py와 계산 로직은 같다(5-fold 체크포인트 각각에
대해 KernelExplainer로 local SHAP을 구하고 평균 - 재학습 없음). 다른 점은
배경표본을 원본 데이터에서 실시간으로 kmeans 요약하지 않고
WEB/scripts/precompute_artifacts.py가 미리 만들어둔 dm_shap_background.csv
(kmeans 합성 대표점 100개, 실제 개인 레코드 아님)를 그대로 쓴다는 것 -
원본 데이터를 Space에 안 올리기 위함이자, 매 요청마다 kmeans를 다시 안
돌려도 되므로 더 빠르다(무료 CPU에서 특히 중요).

체크포인트는 PRS_DM을 넣어 학습한 모델(ML/checkpoints/dm_prs)이다. prs_value
(z 점수)를 주면 PRS를 mask=1로 넣고 PRS_DM도 하나의 SHAP feature로 설명한다.
None이면 PRS를 [MASK]로 두고 EHR feature만 설명한다(학습 때 --prs-dropout으로
[MASK]도 학습돼 있어 EHR만으로 예측 가능).
"""
from pathlib import Path

import numpy as np
import pandas as pd
import shap
import torch

from models.mlp import MultimodalModel
from utils.shap_utils import build_predict_fn


def compute_dm_local_shap(
    patient_row: dict,
    features: list,
    background_df: pd.DataFrame,
    checkpoint_dir: str = "artifacts/checkpoints",
    n_folds: int = 5,
    prs_dim: int = 1,
    device: str = "cpu",
    top_k: int = 8,
    prs_value: float = None,
    prs_feature: str = "PRS_DM",
    exclude_features: tuple = (),
) -> tuple:
    """반환: (importance_df(top_k, feature/patient_value/local_shap/direction),
    ensemble_predicted_prob). 5개 fold 체크포인트 각각으로 이 한 사람의 local
    SHAP을 구하고 feature별 평균을 최종값으로 쓴다(예측 확률도 5개 모델 평균).
    exclude_features: 예측에는 그대로 쓰되 설명 표에서만 뺄 feature - 사용자가
    입력하지 않아 인구 중간값으로 채운 항목은 중간값이 SHAP 배경 평균과 달라
    (예: ALT 중간값 19 vs 평균 22.7) 입력하지도 않은 값에 기여도가 생기므로 뺀다."""
    prs_in_input = prs_value is not None
    ehr_features = list(features)
    patient_row = dict(patient_row)
    if prs_in_input:
        features = ehr_features + [prs_feature]
        patient_row[prs_feature] = float(prs_value)
    patient_vec = np.array([[patient_row[f] for f in features]], dtype=np.float32)
    background = background_df[features].to_numpy(dtype=np.float32)

    shap_values_per_fold, probs = [], []
    for fold_id in range(n_folds):
        checkpoint_path = Path(checkpoint_dir) / f"dm_fold{fold_id}.pt"
        model = MultimodalModel(ehr_dim=len(ehr_features), prs_dim=prs_dim).to(device)
        model.load_state_dict(torch.load(checkpoint_path, map_location=device))
        model.eval()

        predict_fn = build_predict_fn(
            model, prs_dim=max(prs_dim, 1), device=device, apply_sigmoid=True, prs_in_input=prs_in_input
        )
        explainer = shap.KernelExplainer(predict_fn, background)
        sv = np.asarray(explainer.shap_values(patient_vec)).reshape(1, len(features))
        shap_values_per_fold.append(sv[0])
        probs.append(float(predict_fn(patient_vec)[0]))

    mean_shap = np.mean(shap_values_per_fold, axis=0)
    ensemble_prob = float(np.mean(probs))

    df = pd.DataFrame(
        {"feature": features, "patient_value": [patient_row[f] for f in features], "local_shap": mean_shap}
    )
    df = df[~df["feature"].isin(exclude_features)]
    df = df.assign(_abs=df["local_shap"].abs()).sort_values("_abs", ascending=False).drop(columns="_abs")
    df = df.head(top_k).reset_index(drop=True)
    df["direction"] = df["local_shap"].apply(lambda x: "당뇨병 위험 증가" if x > 0 else "당뇨병 위험 감소")
    return df, ensemble_prob