"""
kfold_mlp_app.py — 피하주사 약물동태 예측 앱 (K-Fold 앙상블)
==============================================================
mlp_kfold.py 로 학습한 5개 fold 모델을 모두 불러와
앙상블 평균으로 예측하고, fold 간 편차로 신뢰구간을 표시한다.

사전 준비: python mlp_kfold.py 실행 후 outputs_kfold/folds/ 생성
실행:      streamlit run kfold_mlp_app.py
"""

import matplotlib
import matplotlib.pyplot as plt
import io
import numpy as np
import pandas as pd
import streamlit as st
import tensorflow as tf
import keras
from keras import layers
from pathlib import Path

BASE_DIR    = Path(__file__).resolve().parent
OUTPUT_DIR  = BASE_DIR / "outputs_kfold"
FOLDS_DIR   = OUTPUT_DIR / "folds"

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"}
N_PARAMS  = 10
N_OUTPUTS = 4

# === 모델 구조 (mlp_kfold.py와 반드시 동일) ===
HIDDEN_DIMS = [128, 256, 256, 128]
L2          = 1e-5

PARAM_RANGES = {
    "Lp_ve":    (4e-12,  1.6e-11),
    "K":        (1e-16,  1e-13),
    "p_le":     (1e-9,   1e-7),
    "sigma_ve": (0.01,   0.99),
    "sigma_le": (0.01,   0.50),
    "p_ve":     (1e-11,  1e-8),
    "D_gel":    (1e-12,  1e-9),
    "k_decay":  (1e-8,   1e-4),
    "kf_m":     (0.1,    3.0),
    "kr_m":     (2.5,    200.0),
}
DEFAULT_VALUES = np.array([
    8e-12, 1e-15, 1e-8, 0.9, 0.1,
    1e-9, 45e-12, 1.7e-6, 0.48, 4.2,
])
COLORS = ["#2ecc71", "#e74c3c", "#3498db", "#f39c12"]
_trapz = np.trapezoid if hasattr(np, "trapezoid") else np.trapz


def _setup_font():
    candidates = ["AppleGothic", "Malgun Gothic", "NanumGothic", "Noto Sans CJK KR"]
    available = {f.name for f in matplotlib.font_manager.fontManager.ttflist}
    for name in candidates:
        if name in available:
            matplotlib.rc("font", family=name)
            break
    matplotlib.rc("axes", unicode_minus=False)

_setup_font()


def build_mlp(n_t):
    """mlp_kfold.py의 build_mlp와 동일 구조 (가중치 로드용)."""
    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)
    out = layers.Dense(n_t * N_OUTPUTS, activation="linear", name="curve")(x)
    return keras.Model(inp, out, name="MLP_fold")


@st.cache_resource
def load_folds():
    """folds/fold*/ 의 가중치 + 스케일러를 모두 로드."""
    if not FOLDS_DIR.exists():
        raise FileNotFoundError(
            f"folds 폴더를 찾을 수 없습니다.\n경로: {FOLDS_DIR.resolve()}\n\n"
            "먼저 학습을 실행하세요:\n  python mlp_kfold.py"
        )
    fold_dirs = sorted([d for d in FOLDS_DIR.iterdir()
                        if d.is_dir() and d.name.startswith("fold")])
    if not fold_dirs:
        raise FileNotFoundError(f"fold 모델이 없습니다: {FOLDS_DIR.resolve()}")

    # 시간축은 fold 공통 (첫 fold 기준)
    sc0 = np.load(fold_dirs[0] / "scalers.npz")
    time_arr = sc0["time_arr"]
    n_t = len(time_arr)

    folds = []
    for fd in fold_dirs:
        sc = np.load(fd / "scalers.npz")
        model = build_mlp(n_t)
        model.load_weights(str(fd / "model.weights.h5"))
        folds.append({
            "name": fd.name,
            "model": model,
            "param_mean": sc["param_mean"],
            "param_std":  sc["param_std"],
            "param_log_mask": sc["param_log_mask"].astype(bool),
            "out_mean": sc["out_mean"],
            "out_std":  sc["out_std"],
        })
    return folds, time_arr, n_t


def preprocess(params, fold):
    lm = fold["param_log_mask"]
    pt = params.copy().astype(np.float64)
    pt[lm] = np.log10(np.clip(pt[lm], 1e-300, None))
    return ((pt - fold["param_mean"]) / fold["param_std"]).astype(np.float32)


def predict_ensemble(folds, params, n_t):
    """
    각 fold로 예측 후 스택.
    반환: preds (n_folds, n_t, N_OUTPUTS)
    """
    preds = []
    for fold in folds:
        x = preprocess(params, fold).reshape(1, -1)
        p = fold["model"](x, training=False).numpy().reshape(n_t, N_OUTPUTS)
        p = p * fold["out_std"] + fold["out_mean"]
        preds.append(p)
    return np.stack(preds, axis=0)


def check_range(params):
    warnings = []
    for name, val in zip(PARAM_NAMES, params):
        lo, hi = PARAM_RANGES[name]
        if val < lo or val > hi:
            warnings.append(f"**{name}** = {val:.3e}  (허용: {lo:.1e} ~ {hi:.1e})")
    return warnings


def plot_combined(mean, std, time_arr):
    fig, ax = plt.subplots(figsize=(12, 5))
    for k, name in enumerate(OUTPUT_NAMES):
        ax.plot(time_arr, mean[:, k], color=COLORS[k], lw=2.5, label=name)
        ax.fill_between(time_arr, mean[:, k] - std[:, k], mean[:, k] + std[:, k],
                        color=COLORS[k], alpha=0.18)
    ax.set_xlabel("Time (min)", fontsize=12)
    ax.set_ylabel("농도 (%)", fontsize=12)
    ax.set_title("약물동태 예측", fontsize=13, fontweight="bold")
    ax.set_xlim(time_arr.min(), time_arr.max())
    ax.set_ylim(bottom=0)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    return fig


def plot_subplots(mean, std, time_arr):
    fig, axes = plt.subplots(2, 2, figsize=(14, 7))
    axes = axes.flatten()
    for k, name in enumerate(OUTPUT_NAMES):
        ax = axes[k]
        ax.plot(time_arr, mean[:, k], color=COLORS[k], lw=2.5)
        ax.fill_between(time_arr, mean[:, k] - std[:, k], mean[:, k] + std[:, k],
                        color=COLORS[k], alpha=0.18)
        ax.set_title(name, fontsize=12, fontweight="bold")
        ax.set_xlabel("Time (min)", fontsize=10)
        ax.set_ylabel("농도 (%)", fontsize=10)
        ax.set_xlim(time_arr.min(), time_arr.max())
        ax.set_ylim(bottom=0)
        ax.grid(True, alpha=0.3)
    plt.suptitle("구획별 농도-시간 곡선 (± 편차)", fontsize=13, y=1.01)
    plt.tight_layout()
    return fig


st.set_page_config(page_title="PK Surrogate — K-Fold Ensemble", page_icon="💉", layout="wide")
st.title("💉 피하주사 약물동태 예측 모델")
st.markdown("10개 파라미터 → 4구획 농도-시간 곡선 (0~72시간) | "
            "음영은 fold 간 편차(예측 신뢰도)")

try:
    folds, time_arr, n_t = load_folds()
except FileNotFoundError as e:
    st.error(str(e))
    st.stop()
except Exception as e:
    st.error(f"로딩 실패: {e}")
    st.stop()

st.sidebar.header("📋 파라미터 입력")
st.sidebar.markdown("---")
input_params = []
for i, name in enumerate(PARAM_NAMES):
    dv = float(DEFAULT_VALUES[i])
    lo, hi = PARAM_RANGES[name]
    fmt = "%.2e" if (abs(dv) < 1e-3 or abs(dv) > 1e3) else "%.4f"
    val = st.sidebar.number_input(
        label=f"{name}  ({lo:.1e} ~ {hi:.1e})",
        value=dv, format=fmt, key=f"p_{i}",
    )
    input_params.append(val)
st.sidebar.markdown("---")
predict_btn = st.sidebar.button("🔮 예측하기", type="primary", use_container_width=True)

with st.expander("ℹ️ 모델 정보", expanded=False):
    c1, c2, c3, c4 = st.columns(4)
    c1.metric("모델", "MLP")
    c2.metric("구조", "[128,256,256,128]")
    c3.metric("Dropout", "0.0  (L2=1e-5)")
    c4.metric("Fold 수", f"{len(folds)}개")

if predict_btn:
    params = np.array(input_params, dtype=np.float64)

    oor = check_range(params)
    if oor:
        st.warning(
            "⚠️ **외삽 경고**: 아래 파라미터가 학습 범위를 벗어났습니다. "
            "예측 신뢰도가 낮을 수 있습니다.\n\n"
            + "\n".join(f"- {w}" for w in oor)
        )

    with st.spinner("예측 중..."):
        preds = predict_ensemble(folds, params, n_t)   # (n_folds, n_t, 4)
        mean = preds.mean(axis=0)
        std  = preds.std(axis=0)

    # fold 간 편차가 크면 신뢰도 낮음 경고
    rel_spread = float(std.mean())
    if rel_spread > 1.0:  # 평균 편차 1%p 이상이면
        st.warning(
            f"⚠️ **예측 불확실성 높음**: fold 간 평균 편차가 {rel_spread:.2f}%p입니다. "
            "이 파라미터 영역은 학습 데이터가 성겨 신뢰도가 낮을 수 있습니다. "
            "COMSOL 직접 검증을 권장합니다."
        )

    st.markdown("---")
    st.subheader("📈 72시간 최종값 요약")
    # 총농도(4구획 합)를 맨 앞에 추가
    total_mean = mean.sum(axis=1)   # (n_t,) 시점별 4구획 합
    total_std  = std.sum(axis=1)
    cols = st.columns(5)
    cols[0].metric(
        label="c_total (총합)",
        value=f"{total_mean[-1]:.2f}%",
        delta=f"max: {total_mean.max():.2f}%",
        delta_color="off",
    )
    for k, name in enumerate(OUTPUT_NAMES):
        cols[k + 1].metric(
            label=name,
            value=f"{mean[-1, k]:.2f}%",
            delta=f"±{std[-1, k]:.2f}%p (fold 편차)",
            delta_color="off",
        )

    st.markdown("---")
    st.subheader("📊 농도-시간 곡선")
    tab1, tab2 = st.tabs(["전체 비교", "구획별 상세"])
    with tab1:
        fig1 = plot_combined(mean, std, time_arr)
        st.pyplot(fig1); plt.close(fig1)
    with tab2:
        fig2 = plot_subplots(mean, std, time_arr)
        st.pyplot(fig2); plt.close(fig2)

    st.markdown("---")
    st.subheader("📋 PK 지표")
    # %AUC: 각 구획 AUC가 전체 AUC 합에서 차지하는 비율 (합 100%)
    auc_vals = np.array([_trapz(mean[:, k], time_arr) for k in range(N_OUTPUTS)])
    auc_total = auc_vals.sum() + 1e-12
    rows = []
    for k, name in enumerate(OUTPUT_NAMES):
        curve = mean[:, k]
        rows.append({
            "구획":        name,
            "Cmax (%)":    f"{curve.max():.2f}",
            "Tmax (min)":  f"{time_arr[int(np.argmax(curve))]:.0f}",
            "%AUC":        f"{auc_vals[k] / auc_total * 100:.1f}%",
            "72hr (%)":    f"{curve[-1]:.2f}",
        })
    st.table(rows)

    # ── 시점별 농도값 엑셀 다운로드 ──
    st.markdown("---")
    st.subheader("📥 시점별 농도값 다운로드")
    st.markdown("모델이 예측한 42개 시점의 구획별 농도값을 엑셀로 저장합니다.")

    df_out = pd.DataFrame({"time_min": time_arr, "time_hr": time_arr / 60.0})
    for k, name in enumerate(OUTPUT_NAMES):
        df_out[name] = mean[:, k]

    buf = io.BytesIO()
    with pd.ExcelWriter(buf, engine="openpyxl") as writer:
        df_out.to_excel(writer, index=False, sheet_name="농도_시간별")
        # 입력 파라미터도 별도 시트로 기록 (재현용)
        pd.DataFrame({"parameter": PARAM_NAMES, "value": input_params}
                     ).to_excel(writer, index=False, sheet_name="입력_파라미터")
    buf.seek(0)

    st.download_button(
        label="⬇️ 엑셀 다운로드 (.xlsx)",
        data=buf,
        file_name="농도_시간별_예측.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        use_container_width=False,
    )
    with st.expander("미리보기"):
        st.dataframe(df_out.style.format({
            "time_min": "{:.0f}", "time_hr": "{:.2f}",
            **{name: "{:.3f}" for name in OUTPUT_NAMES},
        }), use_container_width=True)

    with st.expander("🔍 입력 파라미터 확인"):
        st.json({
            name: (f"{val:.3e}" if (abs(val) < 1e-3 or abs(val) > 1e3) else f"{val:.4f}")
            for name, val in zip(PARAM_NAMES, input_params)
        })

else:
    st.info("👈 왼쪽 사이드바에서 파라미터를 입력하고 **예측하기** 버튼을 누르세요.")
    st.markdown("""
    ### 예측 신뢰도
    fold 간 **편차(음영)**가 그 파라미터 영역에서의 예측 신뢰도를 나타냅니다.
    편차가 크면 학습 데이터가 성긴 영역이라 신뢰도가 낮습니다.

    ### 출력 구획
    | 구획 | 설명 |
    |------|------|
    | c_lymph | 림프관 배출 |
    | c_vessel | 혈관 배출 |
    | c_decay | 분해 |
    | c_ecm | 세포외기질 잔류 |

    ### 주의사항
    - 파라미터 범위를 벗어나면 **외삽 경고**
    - fold 편차가 크면 **불확실성 경고** (COMSOL 직접 검증 권장)
    """)

if __name__ == "__main__":
    import os, subprocess, sys
    if os.environ.get("STREAMLIT_RUNNING") != "1":
        env = os.environ.copy()
        env["STREAMLIT_RUNNING"] = "1"
        subprocess.run([sys.executable, "-m", "streamlit", "run", __file__], env=env)