"""
mlp_kfold.py — MLP 단일 모델 + K-Fold 교차검증
================================================
단일 분할의 운(運)을 제거하고 전체 데이터에 대해
평균±std로 신뢰성 있는 성능을 평가한다.

설계:
  - 층화 K-Fold (k_decay × p_ve 기준)로 데이터를 K등분
  - 각 fold가 한 번씩 test가 됨 (전체 케이스가 정확히 1번씩 평가됨)
  - 나머지에서 val 일부 분리, 나머지 train
  - K개 모델의 test 지표를 평균±std로 보고
  - 각 fold 모델/스케일러 저장

모델 설정 (기존 최적):
  구조 [128,256,256,128] SiLU, dropout=0, l2=1e-5
  시간가중 MSE(tau=2000), epochs=1000, patience_es=80

실행: python mlp_kfold.py
"""

from __future__ import annotations
import os
os.environ["PYTHONHASHSEED"] = "42"
os.environ["TF_DETERMINISTIC_OPS"] = "1"
os.environ["TF_CUDNN_DETERMINISTIC"] = "1"

import json
import random
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm

def set_korean_font():
    candidates = ["AppleGothic", "Malgun Gothic", "NanumGothic",
                  "NanumBarunGothic", "DejaVu Sans"]
    available = {f.name for f in fm.fontManager.ttflist}
    for font in candidates:
        if font in available:
            plt.rcParams["font.family"] = font
            break
    plt.rcParams["axes.unicode_minus"] = False  # 유니코드 마이너스 대신 ASCII 하이픈 사용

set_korean_font()

# 폰트 글리프 경고 숨김 (U+2212 등 무해한 치환 경고)
import logging
logging.getLogger("matplotlib.font_manager").setLevel(logging.ERROR)
import warnings
warnings.filterwarnings("ignore", message="Glyph.*missing from font")
warnings.filterwarnings("ignore", message=".*does not have a glyph.*")
from sklearn.metrics import mean_absolute_error, r2_score

import tensorflow as tf
tf.get_logger().setLevel("ERROR")  # retracing 등 무해한 경고 숨김
import keras
from keras import layers, callbacks


# ─────────────────────────────────────────────
# 0. 설정
# ─────────────────────────────────────────────
BASE_DIR   = Path(__file__).resolve().parent
DATA_PATH  = BASE_DIR / "정제_전체_long.xlsx"
OUTPUT_DIR = BASE_DIR / "outputs_kfold"
for d in (OUTPUT_DIR, OUTPUT_DIR / "folds"):
    d.mkdir(parents=True, exist_ok=True)

SEED      = 42
N_OUTPUTS = 4
N_PARAMS  = 10

PARAM_NAMES  = ["Lp_ve", "K", "p_le", "sigma_ve", "sigma_le",
                "p_ve", "D_gel", "k_decay", "kf_m", "kr_m"]
OUTPUT_NAMES = ["c_lymph", "c_vessel", "c_decay", "c_ecm"]
LOG_TRANSFORM_PARAMS = {"K", "p_le", "p_ve", "D_gel", "k_decay", "kr_m"}
EXCLUDE_CASES = {304, 312, 313, 317}

# === K-Fold 설정 ===
N_FOLDS       = 5      # 5-fold (각 fold test ≈ 20%)
VAL_RATIO     = 0.15   # train 내에서 val로 떼는 비율
STRATIFY_BINS = 5      # 층화 기준 분위수

# === 외딴 케이스 제외 설정 ===
# 곡선이 파라미터 이웃과 크게 다른(=데이터가 성겨 보간 불가능한) 케이스를
# 학습/평가 전에 제외. 모델 결과가 아닌 데이터 구조만으로 판정하므로 객관적.
# 목적: 데이터 커버리지가 충분한 정상 영역에서의 본래 성능 측정.
#       제외된 영역은 추가 COMSOL 샘플링으로 별도 보완.
EXCLUDE_SPARSE   = True   # True: 외딴 케이스 제외, False: 전체 사용
SPARSE_NEIGHBORS = 5      # 국소 불일치 계산용 이웃 수
SPARSE_IQR_K     = 1.5    # 이상치 임계 (Q3 + K×IQR)

# === 모델 설정 (기존 최적) ===
HIDDEN_DIMS     = [128, 256, 256, 128]
DROPOUT         = 0.0
L2              = 1e-5
TIME_WEIGHT_TAU = 2000.0
EPOCHS          = 1000
BATCH_SIZE      = 16
LEARNING_RATE   = 1e-3
WEIGHT_DECAY    = 1e-5
PATIENCE_ES     = 80
PATIENCE_LR     = 30
MIN_DELTA       = 1e-5  # 이보다 작은 val_loss 개선은 무시 (조기종료 정상 작동)


def set_seed(seed=SEED):
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)
    try:
        tf.config.experimental.enable_op_determinism()
    except Exception:
        pass


# ─────────────────────────────────────────────
# 1. 데이터 로드
# ─────────────────────────────────────────────
def load_data(path):
    df = pd.read_excel(path)
    required = {"case", "time_min", *PARAM_NAMES, *OUTPUT_NAMES}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"누락된 컬럼: {missing}")

    if EXCLUDE_CASES:
        before = df["case"].nunique()
        df = df[~df["case"].isin(EXCLUDE_CASES)].copy()
        print(f"  제외된 발산 케이스: {sorted(EXCLUDE_CASES)} ({before}->{df['case'].nunique()})")

    cases = sorted(df["case"].unique())
    n = len(cases)
    ref_t = df[df["case"] == cases[0]].sort_values("time_min")["time_min"].values
    n_t = len(ref_t)

    P = np.zeros((n, N_PARAMS), dtype=np.float64)
    C = np.zeros((n, n_t, N_OUTPUTS), dtype=np.float64)
    for i, c in enumerate(cases):
        sub = df[df["case"] == c].sort_values("time_min")
        if len(sub) != n_t or not np.allclose(sub["time_min"].values, ref_t):
            raise ValueError(f"case {c}: 시간 그리드 불일치")
        P[i] = sub[PARAM_NAMES].iloc[0].values
        C[i] = sub[OUTPUT_NAMES].values

    n_clip = (C < 0).sum()
    if n_clip > 0:
        print(f"  음수 농도 클리핑: {n_clip}개 -> 0 (최소 {C.min():.4f})")
        C = np.clip(C, 0.0, None)

    print(f"  로드 완료: 케이스 {n}개, 시간점 {n_t}개")
    return P, C, ref_t, cases


# ─────────────────────────────────────────────
# 2. 스케일러
# ─────────────────────────────────────────────


# ─────────────────────────────────────────────
# 1-b. 외딴(sparse) 케이스 탐지
# ─────────────────────────────────────────────
def detect_sparse_cases(params, curves):
    """
    각 케이스의 곡선이 '파라미터 공간 이웃들의 곡선'과 얼마나 다른지(국소 불일치)
    측정해, IQR 이상치를 외딴 케이스로 판정한다.
    국소 불일치가 크다 = 파라미터는 가까운데 거동이 급변 = 데이터 해상도 부족 영역.

    반환: keep_mask (True=유지), sparse_idx (제외 인덱스), mismatch (점수)
    """
    from scipy.spatial.distance import cdist
    log_mask = np.array([n in LOG_TRANSFORM_PARAMS for n in PARAM_NAMES], dtype=bool)
    Pl = params.copy().astype(float)
    Pl[:, log_mask] = np.log10(np.clip(Pl[:, log_mask], 1e-300, None))
    Psc = (Pl - Pl.mean(0)) / (Pl.std(0) + 1e-12)

    D = cdist(Psc, Psc)
    np.fill_diagonal(D, np.inf)

    mismatch = np.zeros(len(params))
    for i in range(len(params)):
        nn = np.argsort(D[i])[:SPARSE_NEIGHBORS]
        neighbor_mean = curves[nn].mean(axis=0)
        mismatch[i] = np.abs(curves[i] - neighbor_mean).mean()

    q1, q3 = np.percentile(mismatch, [25, 75])
    thr = q3 + SPARSE_IQR_K * (q3 - q1)
    sparse_idx = np.where(mismatch > thr)[0]
    keep_mask = mismatch <= thr
    return keep_mask, sparse_idx, mismatch, thr

class ParamScaler:
    def __init__(self):
        self.log_mask_ = np.array([n in LOG_TRANSFORM_PARAMS for n in PARAM_NAMES], dtype=bool)
        self.mean_ = None; self.std_ = None

    def _apply_log(self, X):
        Xt = X.copy().astype(np.float64)
        Xt[:, self.log_mask_] = np.log10(np.clip(Xt[:, self.log_mask_], 1e-300, None))
        return Xt

    def fit(self, X):
        Xt = self._apply_log(X)
        self.mean_ = Xt.mean(0); self.std_ = Xt.std(0) + 1e-12
        return self

    def transform(self, X):     return (self._apply_log(X) - self.mean_) / self.std_
    def fit_transform(self, X): return self.fit(X).transform(X)


class OutputScaler:
    def __init__(self):
        self.mean_ = None; self.std_ = None

    def fit(self, Y):
        Yf = Y.reshape(-1, N_OUTPUTS)
        self.mean_ = Yf.mean(0); self.std_ = Yf.std(0) + 1e-12
        return self

    def transform(self, Y):         return (Y - self.mean_) / self.std_
    def inverse_transform(self, Y): return Y * self.std_ + self.mean_
    def fit_transform(self, Y):     return self.fit(Y).transform(Y)


# ─────────────────────────────────────────────
# 3. 층화 K-Fold 분할
# ─────────────────────────────────────────────
def make_strata(params):
    """k_decay × p_ve 기준 층 라벨 생성."""
    kd = np.log10(np.clip(params[:, PARAM_NAMES.index("k_decay")], 1e-300, None))
    pv = np.log10(np.clip(params[:, PARAM_NAMES.index("p_ve")], 1e-300, None))
    kb = pd.qcut(kd, q=STRATIFY_BINS, labels=False, duplicates="drop")
    pb = pd.qcut(pv, q=STRATIFY_BINS, labels=False, duplicates="drop")
    return kb * (STRATIFY_BINS + 1) + pb


def stratified_kfold_indices(params, n_folds, seed):
    """
    층화 K-Fold: 각 층 내부를 n_folds로 나눠 fold마다 고르게 분배.
    반환: fold_assign (각 케이스의 fold 번호 0..n_folds-1)
    """
    rng = np.random.default_rng(seed)
    strata = make_strata(params)
    fold_assign = np.full(len(params), -1, dtype=int)
    for s in np.unique(strata):
        idx = np.where(strata == s)[0]
        rng.shuffle(idx)
        # 층 내부를 순환 배정 (각 fold에 고르게)
        for j, i in enumerate(idx):
            fold_assign[i] = j % n_folds
    return fold_assign


def split_train_val(train_idx, params, val_ratio, seed):
    """train 인덱스에서 val을 층화로 분리."""
    rng = np.random.default_rng(seed)
    strata = make_strata(params)[train_idx]
    tr, va = [], []
    for s in np.unique(strata):
        local = np.where(strata == s)[0]
        rng.shuffle(local)
        nv = max(1, round(len(local) * val_ratio))
        if len(local) - nv < 1:
            tr.extend(train_idx[local].tolist()); continue
        va.extend(train_idx[local[:nv]].tolist())
        tr.extend(train_idx[local[nv:]].tolist())
    return np.array(tr), np.array(va)


# ─────────────────────────────────────────────
# 4. 손실 / 모델
# ─────────────────────────────────────────────
def make_time_weighted_loss(time_arr, n_t):
    w = np.exp(-time_arr / TIME_WEIGHT_TAU)
    w = (w / w.mean()).astype(np.float32)
    tw = tf.constant(w.reshape(1, n_t, 1), dtype=tf.float32)

    @tf.function
    def loss_fn(y_true, y_pred):
        yt = tf.reshape(y_true, (-1, n_t, N_OUTPUTS))
        yp = tf.reshape(y_pred, (-1, n_t, N_OUTPUTS))
        return tf.reduce_mean(tf.square(yt - yp) * tw)
    return loss_fn


def build_mlp(n_t):
    inp = keras.Input(shape=(N_PARAMS,), name="params")
    x = inp
    for h in HIDDEN_DIMS:
        x = layers.Dense(h, kernel_regularizer=keras.regularizers.l2(L2) if L2 > 0 else None)(x)
        x = layers.Activation(tf.nn.silu)(x)
        if DROPOUT > 0:
            x = layers.Dropout(DROPOUT)(x)
    out = layers.Dense(n_t * N_OUTPUTS, activation="linear", name="curve")(x)
    return keras.Model(inp, out, name="MLP_single")


# ─────────────────────────────────────────────
# 5. 평가 지표
# ─────────────────────────────────────────────
def evaluate(y_true, y_pred):
    case_mae = np.abs(y_true - y_pred).mean(axis=(1, 2))
    per = {name: float(mean_absolute_error(y_true[..., k].reshape(-1),
                                           y_pred[..., k].reshape(-1)))
           for k, name in enumerate(OUTPUT_NAMES)}
    return {
        "MAE_overall": float(mean_absolute_error(y_true.reshape(-1), y_pred.reshape(-1))),
        "R2_overall":  float(r2_score(y_true.reshape(-1), y_pred.reshape(-1))),
        "case_MAE_mean": float(case_mae.mean()),
        "case_MAE_max":  float(case_mae.max()),
        "case_MAE_p90":  float(np.percentile(case_mae, 90)),
        "case_MAE_std":  float(case_mae.std()),
        "per_compartment_MAE": per,
    }


# ─────────────────────────────────────────────
# 6. 단일 fold 학습
# ─────────────────────────────────────────────
def train_one_fold(fold, itr, iva, ite, params, curves, time_arr, n_t):
    print(f"\n{'='*55}")
    print(f"  Fold {fold+1}/{N_FOLDS}  (train={len(itr)}, val={len(iva)}, test={len(ite)})")
    print(f"{'='*55}")

    ps = ParamScaler()
    X_tr = ps.fit_transform(params[itr]).astype(np.float32)
    X_va = ps.transform(params[iva]).astype(np.float32)
    X_te = ps.transform(params[ite]).astype(np.float32)

    osc = OutputScaler(); osc.fit(curves[itr])
    Y_tr = osc.transform(curves[itr]).reshape(len(itr), -1).astype(np.float32)
    Y_va = osc.transform(curves[iva]).reshape(len(iva), -1).astype(np.float32)
    Y_te_raw = curves[ite]

    loss_fn = make_time_weighted_loss(time_arr, n_t)

    set_seed(SEED)  # fold마다 동일 초기화 (모델 변동 통제, 분할 효과만 측정)
    model = build_mlp(n_t)
    model.compile(optimizer=keras.optimizers.Adam(LEARNING_RATE, weight_decay=WEIGHT_DECAY),
                  loss=loss_fn, metrics=["mae"])
    cbs = [
        callbacks.EarlyStopping(monitor="val_loss", patience=PATIENCE_ES,
                                min_delta=MIN_DELTA, restore_best_weights=True, verbose=0),
        callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.5,
                                    patience=PATIENCE_LR, min_delta=MIN_DELTA,
                                    min_lr=1e-7, verbose=0),
    ]
    h = model.fit(X_tr, Y_tr, validation_data=(X_va, Y_va),
                  epochs=EPOCHS, batch_size=BATCH_SIZE, callbacks=cbs, verbose=0)
    ep = len(h.history["val_loss"])

    pred = osc.inverse_transform(model.predict(X_te, verbose=0).reshape(-1, n_t, N_OUTPUTS))
    metrics = evaluate(Y_te_raw, pred)
    print(f"  ep={ep}  test MAE 평균={metrics['case_MAE_mean']:.3f}%  "
          f"최대={metrics['case_MAE_max']:.3f}%  R2={metrics['R2_overall']:.4f}")

    # fold 산출물 저장 (가중치 + 스케일러 → 재현/추론용)
    fdir = OUTPUT_DIR / "folds" / f"fold{fold+1}"
    fdir.mkdir(parents=True, exist_ok=True)
    model.save_weights(str(fdir / "model.weights.h5"))
    np.savez(fdir / "scalers.npz",
             param_mean=ps.mean_, param_std=ps.std_, param_log_mask=ps.log_mask_,
             out_mean=osc.mean_, out_std=osc.std_, time_arr=time_arr,
             test_idx=ite)
    return metrics, h, pred, Y_te_raw, ite


# ─────────────────────────────────────────────
# 7. 종합 시각화
# ─────────────────────────────────────────────
def plot_kfold_summary(fold_metrics, histories, save_path):
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))
    fig.suptitle("K-Fold 교차검증 결과", fontsize=14, fontweight="bold", y=1.02)

    # 왼쪽: fold별 평균/최대 MAE 막대
    ax = axes[0]
    folds = [f"F{i+1}" for i in range(len(fold_metrics))]
    means = [m["case_MAE_mean"] for m in fold_metrics]
    maxes = [m["case_MAE_max"] for m in fold_metrics]
    x = np.arange(len(folds)); width = 0.38
    ax.bar(x - width/2, means, width, label="평균 MAE", color="#2E5C8A")
    ax.bar(x + width/2, maxes, width, label="최대 MAE", color="#C62828")
    ax.axhline(np.mean(means), color="#2E5C8A", linestyle="--", linewidth=1,
               alpha=0.7, label=f"평균의 평균 {np.mean(means):.3f}%")
    ax.set_xticks(x); ax.set_xticklabels(folds)
    ax.set_ylabel("MAE (%)", fontsize=11)
    ax.set_title("Fold별 test MAE", fontsize=12)
    ax.legend(fontsize=9); ax.grid(True, alpha=0.3, axis="y")

    # 오른쪽: fold별 val loss 곡선
    ax2 = axes[1]
    for i, h in enumerate(histories):
        vl = h.history["val_loss"]
        ax2.plot(np.arange(1, len(vl)+1), vl, linewidth=1.3, alpha=0.8, label=f"Fold {i+1}")
    ax2.set_xlabel("Epoch", fontsize=11); ax2.set_ylabel("Val Loss", fontsize=11)
    ax2.set_title("Fold별 검증 손실", fontsize=12)
    ax2.set_yscale("log")
    ax2.legend(fontsize=8); ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"\n  요약 그래프 저장: {save_path}")


# ─────────────────────────────────────────────
# 8. 메인
# ─────────────────────────────────────────────
def main():
    print("=" * 60)
    print(f"  MLP {N_FOLDS}-Fold 교차검증")
    print("=" * 60)

    print("\n[1] 데이터 로드")
    params, curves, time_arr, cases = load_data(DATA_PATH)
    n_t = len(time_arr)

    if EXCLUDE_SPARSE:
        print("\n[1-b] 외딴 케이스 제외 (국소 불일치 IQR 기준)")
        keep, sparse_idx, mismatch, thr = detect_sparse_cases(params, curves)
        excl_cases = [cases[i] for i in sparse_idx]
        print(f"  임계값(Q3+{SPARSE_IQR_K}×IQR) = {thr:.2f}%")
        print(f"  제외 {len(sparse_idx)}개 / 유지 {keep.sum()}개")
        print(f"  제외 케이스: {sorted(excl_cases)}")
        params = params[keep]
        curves = curves[keep]
        cases = [c for i, c in enumerate(cases) if keep[i]]

    print(f"\n[2] 층화 {N_FOLDS}-Fold 분할")
    fold_assign = stratified_kfold_indices(params, N_FOLDS, SEED)
    for f in range(N_FOLDS):
        print(f"  Fold {f+1}: test {int((fold_assign==f).sum())}개")

    print("\n[3] Fold별 학습")
    fold_metrics, histories = [], []
    # 전체 케이스에 대한 out-of-fold 예측 (각 케이스가 test였을 때의 예측)
    oof_pred = np.zeros_like(curves)
    oof_filled = np.zeros(len(cases), dtype=bool)

    for f in range(N_FOLDS):
        ite = np.where(fold_assign == f)[0]
        rest = np.where(fold_assign != f)[0]
        itr, iva = split_train_val(rest, params, VAL_RATIO, SEED + f)
        metrics, h, pred, _, test_idx = train_one_fold(
            f, itr, iva, ite, params, curves, time_arr, n_t)
        fold_metrics.append(metrics); histories.append(h)
        oof_pred[test_idx] = pred
        oof_filled[test_idx] = True

    # ── 종합 ──
    print("\n" + "=" * 60)
    print("  K-Fold 종합 결과")
    print("=" * 60)
    keys = ["case_MAE_mean", "case_MAE_max", "case_MAE_p90", "R2_overall"]
    labels = {"case_MAE_mean": "평균 MAE", "case_MAE_max": "최대 MAE",
              "case_MAE_p90": "p90 MAE", "R2_overall": "R²"}
    summary = {}
    for k in keys:
        vals = np.array([m[k] for m in fold_metrics])
        summary[k] = {"mean": float(vals.mean()), "std": float(vals.std()),
                      "min": float(vals.min()), "max": float(vals.max())}
        unit = "" if k == "R2_overall" else "%"
        print(f"  {labels[k]:10s}: {vals.mean():.3f}{unit} ± {vals.std():.3f}  "
              f"[{vals.min():.3f}, {vals.max():.3f}]")

    # 전체 OOF 지표 (모든 케이스가 정확히 1번씩 test됨 → 가장 신뢰도 높은 단일 수치)
    assert oof_filled.all(), "일부 케이스가 평가되지 않음"
    oof_metrics = evaluate(curves, oof_pred)
    print(f"\n  [전체 OOF] 평균 MAE={oof_metrics['case_MAE_mean']:.3f}%  "
          f"최대={oof_metrics['case_MAE_max']:.3f}%  "
          f"p90={oof_metrics['case_MAE_p90']:.3f}%  R²={oof_metrics['R2_overall']:.4f}")
    print("  [전체 OOF 구획별 MAE]")
    for name, v in oof_metrics["per_compartment_MAE"].items():
        print(f"    {name:<10}: {v:.3f}%")

    print("\n[4] 저장")
    with open(OUTPUT_DIR / "kfold_metrics.json", "w", encoding="utf-8") as fp:
        json.dump({"per_fold": fold_metrics, "summary": summary,
                   "oof": oof_metrics}, fp, indent=2, ensure_ascii=False)
    np.savez(OUTPUT_DIR / "oof_predictions.npz",
             y_true=curves, y_pred=oof_pred, fold_assign=fold_assign, time_arr=time_arr)
    plot_kfold_summary(fold_metrics, histories, OUTPUT_DIR / "kfold_summary.png")
    print(f"  저장: {OUTPUT_DIR}/ (kfold_metrics.json, oof_predictions.npz, "
          f"kfold_summary.png, folds/)")
    print("\n완료.")


if __name__ == "__main__":
    main()