#!/usr/bin/env python3
"""WEB(HF Spaces) 배포용 요약 산출물을 미리 계산해서 WEB/artifacts/에 저장한다.

이 스크립트는 원본 EHR 데이터(data/processed/EHR_BASE_merged.csv 등)에
접근할 수 있는 이 저장소 환경(로컬)에서만 실행한다. **HF Spaces에는 이
스크립트를 올리지 않고 이 스크립트가 만든 WEB/artifacts/ 결과물만 올린다**
- 원본 개인 EHR 레코드를 공개 웹에 배포하지 않기 위함이다. 산출물은 전부
아래 세 종류뿐이고, 어느 것도 개인 레코드 원본이 아니다:
    1. feature별 평균/중간값(집계 통계 - 몇 명이 어떤 값인지는 알 수 없음)
    2. DM SHAP용 배경표본 - shap.kmeans로 10,000명을 100개의 "합성 대표점"
       으로 압축한 것(각 대표점은 실존하는 특정 개인이 아니라 군집 중심 좌표)
    3. Regression/main_regression.py가 이미 적합해둔 OLS 계수,
       ML/main_dm.py가 이미 학습해둔 DM 체크포인트(둘 다 모델 파라미터일
       뿐 원본 데이터가 아님 - 체크포인트 5개 합쳐도 580KB로 아주 작음)

BAG는 lite 버전(ML/config/base_prs_htn.yaml)을 쓴다: 앞선 EHR 15개 회귀에서
FDR 유의했던 5개(SEX/HTN/HTNAG/DM/SBP)에 고혈압 PRS(PRS_HTN)를 covariate로
넣은 Regression/output_prs_htn/effect_sizes.csv(R² 0.1177 - EHR 15개 0.1186과
거의 같으면서 입력은 6개뿐). DM은 PRS 버전(ML/config/dm_prs.yaml) 체크포인트로 PRS_DM을 넣고
--prs-dropout 0.3으로 학습한 ML/checkpoints/dm_prs/(PRS를 안 넣으면 [MASK]로
예측 - 그 [MASK] 임베딩도 학습돼 있어서 EHR만 입력해도 동작함). feature
통계표에 PRS 행(평균/중간값, z-표준화라 ~0)을, DM SHAP 배경점에 PRS_DM
컬럼을 같이 넣는다(PRS를 입력한 사람은 PRS도 SHAP feature로 설명).

WEB 쪽 코드(WEB/utils/patient_data.py 등)는 이 산출물만 읽고, MultimodalDataset
이나 원본 CSV/텍스트 파일을 전혀 참조하지 않는다 - T_ID로 기존 사람을
조회하는 기능도 그래서 WEB에는 없다(임의 EHR 값 입력만 지원).

사용법 (torch/shap/pandas가 있는 hugging_face 환경 필요, 저장소 루트에서 실행):
    /data1/software/anaconda3/envs/hugging_face/bin/python WEB/scripts/precompute_artifacts.py
"""
import shutil
import sys
from pathlib import Path

import pandas as pd
import shap

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO_ROOT))

from src.data_loaders import MultimodalDataset, load_config

ARTIFACTS_DIR = REPO_ROOT / "WEB" / "artifacts"


def save_feature_stats(config_path: str, out_name: str) -> tuple:
    """BASE 전체(fold 없음)에서 feature별 평균/중간값만 뽑아 저장한다.
    반환: (features, 그 df) - df는 이 함수를 부른 뒤 SHAP 배경 계산에도
    재사용하고 함수 밖으로는 안 나감(원본은 여기서만 잠깐 메모리에 있다가 버려짐)."""
    config = load_config(config_path)
    df = MultimodalDataset(config, cohort="base", folds=None).df
    # EHR feature 뒤에 PRS feature(PRS_HTN/PRS_DM)도 붙인다 - WEB/app.py는
    # config.yaml의 prs 설정을 보고 PRS 행을 EHR 입력 목록에서 따로 뺀다.
    features = config["data"]["ehr"]["features"] + list(config["data"]["prs"]["features"])

    stats = pd.DataFrame(
        {
            "feature": features,
            "mean": df[features].mean().to_numpy(),
            "median": df[features].median().to_numpy(),
        }
    )
    out_path = ARTIFACTS_DIR / out_name
    stats.to_csv(out_path, index=False)
    print(f"[INFO] {out_name} 저장 ({len(features)}개 feature 평균/중간값): {out_path}")
    return features, df


def main():
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
    (ARTIFACTS_DIR / "checkpoints").mkdir(exist_ok=True)

    print("[INFO] BAG feature 통계 계산 중 (ML/config/base_prs_htn.yaml)...")
    save_feature_stats("ML/config/base_prs_htn.yaml", "bag_feature_stats.csv")

    print("[INFO] DM feature 통계 계산 중 (ML/config/dm_prs.yaml)...")
    dm_features, dm_df = save_feature_stats("ML/config/dm_prs.yaml", "dm_feature_stats.csv")

    print("[INFO] DM local SHAP 배경표본(kmeans 100개 합성 대표점) 계산 중...")
    background = shap.kmeans(dm_df[dm_features].to_numpy(dtype="float32"), 100)
    background_df = pd.DataFrame(background.data, columns=dm_features)
    background_path = ARTIFACTS_DIR / "dm_shap_background.csv"
    background_df.to_csv(background_path, index=False)
    print(f"[INFO] dm_shap_background.csv 저장 ({len(background_df)}개 합성 대표점, 실제 개인 아님): {background_path}")

    print("[INFO] effect_sizes.csv / DM 5-fold 체크포인트 복사 중...")
    shutil.copy(REPO_ROOT / "Regression/output_prs_htn/effect_sizes.csv", ARTIFACTS_DIR / "effect_sizes.csv")
    for fold_id in range(5):
        src = REPO_ROOT / "ML/checkpoints/dm_prs" / f"dm_fold{fold_id}.pt"
        shutil.copy(src, ARTIFACTS_DIR / "checkpoints" / f"dm_fold{fold_id}.pt")

    total_size = sum(f.stat().st_size for f in ARTIFACTS_DIR.rglob("*") if f.is_file())
    print(f"\n[INFO] 완료 - {ARTIFACTS_DIR} 전체 크기 {total_size / 1024:.1f} KB")
    print("[INFO] 이 폴더 안의 파일만 WEB 배포에 포함하면 됨 (원본 EHR 데이터 없음)")


if __name__ == "__main__":
    main()