# -*- coding: utf-8 -*-
"""산업 VARX 추정·전망 엔진 — 결합 가중치 검증과 상류 재추정(FULL=1)에만 쓴다.

출처(원 프로젝트, 정리 사본):
  industry_pipeline/run_grid.py      load · ar_sigma · design · prior_var · fit_eqwise (2026-08-28판)
  industry_pipeline/run_forecast.py  fit_full · forecast_one · q_to_annual · run (2026-08-28판)
  scratchpad/validate_shrink_w.py    결합 가중치 w의 유사표본외 선택 (2026-09-03판,
                                     통합 번들 사본 industry/code/validate_shrink_w.py와 동일)

원 코드와 달라진 점은 다음 세 가지뿐이며 계산식은 같다.
  1. 경로를 config_env 대신 인수로 받는다(패키지 상대경로).
  2. 결과를 파일에 쓰지 않고 DataFrame·dict로 돌려준다(저장은 노트북이 한다).
  3. 화면 출력을 없앴다.

명세: 차분(y = 100·Δlog VA) · 산업연관표(IO) 사전 · VARX(외생 6종, 동시+2시차), 내생 4시차, 방정식별 베이지언 추정.
총량(탑다운) 모형만 수준(y = 100·log VA, 자기계수 사전평균 1) 명세를 쓴다.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

P, Q = 4, 2            # 내생 시차, 외생 시차(동시 포함 0..Q)
LAM = 0.2              # 미네소타 전체 강도
EXOG = ["유가", "세계수요", "실질금리", "실질환율", "고용_취업자", "우리경기"]
SCENARIOS = ["S1", "S2", "S3", "S4", "S5"]
HIST_Y = list(range(2000, 2026))
FUT_Y = list(range(2026, 2036))


def quarters(start: str = "2000Q1", end: str = "2025Q4") -> list[str]:
    y0, q0 = int(start[:4]), int(start[5]); y1, q1 = int(end[:4]), int(end[5])
    out, y, q = [], y0, q0
    while (y, q) <= (y1, q1):
        out.append(f"{y}Q{q}")
        q += 1
        if q > 4:
            q, y = 1, y + 1
    return out


FUT = quarters("2026Q1", "2035Q4")


# ── run_grid.py ─────────────────────────────────────────────────────────
def load(levels_csv: str, exog_hist_csv: str, io_csv: str | None, gdp_col: str = "총량GDP"):
    """분기 수준(십억원) → L = 100·log VA, 외생 과거값, IO 연관행렬(산업 순서 정렬)."""
    lev = pd.read_csv(levels_csv, index_col=0)
    qs = quarters()
    inds = [c for c in lev.columns if c != gdp_col]
    L = 100.0 * np.log(lev[inds].reindex(qs))
    gdp = 100.0 * np.log(lev[gdp_col].reindex(qs)) if gdp_col in lev.columns else None
    exo = pd.read_csv(exog_hist_csv).set_index("분기").reindex(qs)[EXOG]
    Wio = None
    if io_csv is not None:
        w = pd.read_csv(io_csv, index_col=0)
        common = [c for c in inds if c in w.columns and c in w.index]
        if len(common) == len(inds):
            Wio = w.loc[inds, inds]
        elif len(common) >= 0.5 * len(inds):
            Wio = w.reindex(index=inds, columns=inds).fillna(0.05)   # 일부 미부여는 최소수축
    return L, gdp, exo, Wio


def ar_sigma(Y: pd.DataFrame) -> np.ndarray:
    """단변량 AR(P) 잔차 표준편차(미네소타 척도)."""
    s = np.zeros(Y.shape[1]); V = Y.values
    for i in range(Y.shape[1]):
        y = V[:, i]; yy = y[P:]
        Z = np.column_stack([y[P - l:-l] for l in range(1, P + 1)] + [np.ones(len(yy))])
        b, *_ = np.linalg.lstsq(Z, yy, rcond=None)
        s[i] = (yy - Z @ b).std(ddof=Z.shape[1]) or 1.0
    return np.where(s > 1e-6, s, 1.0)


def design(Y: pd.DataFrame, exo: pd.DataFrame | None, use_x: bool, ect=None):
    idx = list(Y.index); n = Y.shape[1]
    rows = list(range(max(P, Q), len(idx)))
    Ym = Y.values[rows]; cols = []; groups = []
    for l in range(1, P + 1):
        cols.append(Y.values[[r - l for r in rows]]); groups += [("endog", l, j) for j in range(n)]
    if ect is not None:
        cols.append(ect.values[[r - 1 for r in rows]]); groups += [("ect", 1, j) for j in range(n)]
    if use_x and exo is not None:
        for s in range(0, Q + 1):
            cols.append(exo.values[[r - s for r in rows]]); groups += [("exog", s, j) for j in range(exo.shape[1])]
    cols.append(np.ones((len(rows), 1))); groups += [("const", 0, 0)]
    return Ym, np.hstack(cols), groups


def prior_var(groups, sigma, n: int, prior: str, Wio):
    """사전분산: 자기 λ²/l², 교차 (λ²/l²)(σ_i²/σ_j²)[×(|W_ij|+0.05), IO 사전], 외생 4λ², 상수 1e6."""
    k = len(groups); V = np.ones((k, n)); Wm = Wio.values if Wio is not None else None
    for c, (kind, l, j) in enumerate(groups):
        for i in range(n):
            if kind == "endog":
                if i == j:
                    base = (LAM ** 2) / (l ** 2)
                else:
                    base = (LAM ** 2) / (l ** 2) * (sigma[i] ** 2 / sigma[j] ** 2)
                    if prior == "io" and Wm is not None:
                        base *= (abs(Wm[i, j]) + 0.05)
                V[c, i] = base
            elif kind == "ect":
                V[c, i] = LAM ** 2
            elif kind == "exog":
                V[c, i] = (LAM ** 2) * 4.0
            else:
                V[c, i] = 1e6
    return V


def fit_eqwise(Ym, X, V, sigma, delta1: bool):
    """방정식별 정규 사전 사후평균: (V0⁻¹ + X'X/s²)⁻¹(V0⁻¹b0 + X'y/s²)."""
    k = X.shape[1]; n = Ym.shape[1]; XtX = X.T @ X; B = np.zeros((k, n))
    for i in range(n):
        s2 = sigma[i] ** 2
        b0 = np.zeros(k)
        if delta1:
            b0[i] = 1.0
        iV0 = np.diag(1.0 / V[:, i])
        B[:, i] = np.linalg.solve(iV0 + XtX / s2, iV0 @ b0 + (X.T @ Ym[:, i]) / s2)
    return B


# ── run_forecast.py ─────────────────────────────────────────────────────
def fit_full(levels_csv, exog_hist_csv, io_csv, spec="차분", prior="io", em="varx"):
    L, gdp, exo, Wio = load(levels_csv, exog_hist_csv, io_csv)
    n = L.shape[1]; use_x = (em == "varx")
    Y = L.diff() if spec != "수준" else L
    sig = ar_sigma(Y.dropna())
    ect = L.subtract(gdp, axis=0) if (spec == "VECM" and gdp is not None) else None
    Ysub = Y.dropna()
    ecs = ect.loc[Ysub.index] if ect is not None else None
    exs = exo.loc[Ysub.index] if use_x else None
    Ym, X, groups = design(Ysub, exs, use_x, ecs)
    V = prior_var(groups, sig, n, prior, Wio)
    B = fit_eqwise(Ym, X, V, sig, delta1=(spec == "수준"))
    return L, gdp, exo, B, use_x


def forecast_one(L, gdp, exo, B, use_x, spec, exo_future):
    """1기-앞 재귀 전망 2026Q1–2035Q4 → 분기 실질VA(십억원)."""
    inds = list(L.columns)
    efull = pd.concat([exo[EXOG], exo_future[EXOG]])
    Ld = L.diff()
    Lv = {q: L.loc[q].values.astype(float) for q in L.index}
    Gv = {q: Ld.loc[q].values.astype(float) for q in L.index}
    allq = list(L.index) + FUT
    for t in FUT:
        ti = allq.index(t); cols = []
        src = Lv if spec == "수준" else Gv
        for l in range(1, P + 1):
            cols.append(src[allq[ti - l]])
        if spec == "VECM":
            qprev = allq[ti - 1]
            g = gdp.loc[qprev] if (gdp is not None and qprev in gdp.index) else float(gdp.iloc[-1]) if gdp is not None else 0.0
            cols.append(Lv[qprev] - g)
        if use_x:
            for s in range(0, Q + 1):
                cols.append(efull.loc[allq[ti - s]].values.astype(float))
        cols.append(np.array([1.0]))
        yhat = np.concatenate([np.atleast_1d(c) for c in cols]) @ B
        if spec == "수준":
            Lv[t] = yhat; Gv[t] = yhat - Lv[allq[ti - 1]]
        else:
            Lv[t] = Lv[allq[ti - 1]] + yhat; Gv[t] = yhat
    return pd.DataFrame({t: np.exp(Lv[t] / 100.0) for t in FUT}, index=inds).T


def q_to_annual(df_q: pd.DataFrame) -> pd.DataFrame:
    yr = pd.Series([int(q[:4]) for q in df_q.index], index=df_q.index)
    return df_q.groupby(yr).sum()


def sum_growth_mean(L: pd.DataFrame) -> float:
    """표본(2000Q1–2025Q4) 산업 합계의 평균 성장률(연율 %) = 400 × 평균 Δln(Σ_i VA_i)."""
    S = np.exp(L / 100.0).sum(axis=1)
    return float(400.0 * np.log(S).diff().mean())


def anchor_drift(L: pd.DataFrame, B: np.ndarray, target: float) -> tuple[np.ndarray, float]:
    """드리프트 재설정(절편 보정, 2026-10 정비): 차분 VARX의 장기 성장률을 모든 산업에서 같은 폭 δ만큼 낮춘다.

    차분 VARX Δy = c + Σ_l A_l Δy_{t−l} + Γx 의 장기 평균은 μ = (I − ΣA)⁻¹(c + Γx̄)이며, 표본 안에서는 산업별 평균 성장률이다.
    δ = (표본 산업 합계 평균 성장률 − target)/4 (분기 100·log)로 두고 c ← c − (I − ΣA)·δ·1 을 쓰면 μ가 모든 산업에서 δ씩 낮아진다.
    target은 3장 BOKDPM의 잠재성장률 정상상태(growth_kor_ss)이다. 산업 간 성장 격차와 외생·동학 계수는 바꾸지 않는다.
    반환: (조정한 계수행렬, 연율 δ(%p))
    """
    n = L.shape[1]
    d_ann = sum_growth_mean(L) - target
    lagsum = B[:P * n, :].reshape(P, n, n).sum(axis=(0, 1))       # 방정식 i: Σ_l Σ_j A_l[i, j]
    B2 = B.copy()
    B2[-1, :] = B[-1, :] - (1.0 - lagsum) * d_ann / 4.0
    return B2, d_ann


def run_forecast(levels_csv, exog_hist_csv, exog_scen_csv, io_csv, spec="차분", prior="io", em="varx",
                 anchor: float | None = None) -> pd.DataFrame:
    """원 run_forecast.run과 같은 장형 표(시나리오·연도·산업·실질VA_십억원)를 돌려준다.

    anchor(연 %)를 주면 차분 명세의 드리프트를 재설정한다(anchor_drift). None이면 원 계산과 같다.
    """
    L, gdp, exo, B, use_x = fit_full(levels_csv, exog_hist_csv, io_csv, spec, prior, em)
    if anchor is not None:
        assert spec == "차분", "드리프트 재설정은 차분 명세에만 정의한다"
        B, _ = anchor_drift(L, B, anchor)
    inds = list(L.columns)
    hist_q = pd.DataFrame({q: np.exp(L.loc[q].values / 100.0) for q in L.index}, index=inds).T
    hist_a = q_to_annual(hist_q)
    scen = pd.read_csv(exog_scen_csv, encoding="utf-8-sig")
    scen.columns = [c.strip() for c in scen.columns]
    rows = []
    for y in HIST_Y:
        for ind in inds:
            rows.append(("실적", y, ind, round(float(hist_a.loc[y, ind]), 3)))
    for s in pd.unique(scen["시나리오"]):
        ex_s = scen[scen["시나리오"] == s].set_index("분기")
        fc_a = q_to_annual(forecast_one(L, gdp, exo, B, use_x, spec, ex_s))
        for y in FUT_Y:
            for ind in inds:
                rows.append((s, y, ind, round(float(fc_a.loc[y, ind]), 3)))
    return pd.DataFrame(rows, columns=["시나리오", "연도", "산업", "실질VA_십억원"])


def scenario_exog_decomp(levels_csv, exog_hist_csv, exog_scen_csv, io_csv, spec="차분", prior="io", em="varx",
                         anchor: float | None = None, year: int = 2035) -> pd.DataFrame:
    """시나리오 차이(각 시나리오 − S1, year년 합계 %)를 외생별로 나눈다(2026-10 정비판 진단).

    외생 k의 기여 = S1 외생 경로에서 k 하나만 시나리오 s의 경로로 바꾼 전망 − S1 전망. 선형이 아니므로(로그 수준 합) 기여의 합은 전체 차이와 조금 다르다.
    """
    L, gdp, exo, B, use_x = fit_full(levels_csv, exog_hist_csv, io_csv, spec, prior, em)
    if anchor is not None:
        B, _ = anchor_drift(L, B, anchor)
    scen = pd.read_csv(exog_scen_csv, encoding="utf-8-sig")
    ex = {s: scen[scen["시나리오"] == s].set_index("분기") for s in SCENARIOS}
    tot = lambda e: float(q_to_annual(forecast_one(L, gdp, exo, B, use_x, spec, e)).sum(axis=1).loc[year])  # noqa: E731
    b = tot(ex["S1"])
    rows = []
    for s in SCENARIOS[1:]:
        r = {"시나리오": s, "전체(%)": 100 * (tot(ex[s]) / b - 1)}
        for k in EXOG:
            e = ex["S1"].copy()
            e[k] = ex[s][k]
            r[f"{k}(%)"] = 100 * (tot(e) / b - 1)
        rows.append(r)
    return pd.DataFrame(rows).set_index("시나리오")


# ── validate_shrink_w.py ────────────────────────────────────────────────
def validate_shrink_w(levels_csv, exog_hist_csv, io_csv, train_end="2019Q4",
                      val_years=(2020, 2021, 2022), test_years=(2023, 2024, 2025)) -> dict:
    """기존 전망 결합 가중치 w의 유사표본외 선택(62분류).

    - 모형: 차분·IO·VARX를 2000Q1–train_end로 적합, 2020Q1–2025Q4를 실측 외생으로 1기-앞 재귀 전망.
    - 기존 전망의 과거 대용: 각 산업 로그수준을 학습표본 평균 로그성장으로 연장한 추세 전망.
    - 결합: ỹ = w·추세 + (1−w)·모형(로그수준). 연간 수준(십억원)으로 환산해 산업 평균제곱오차의 제곱근.
    - w* = 검증(2020–22) 최소, 시험(2023–25)에서 성능 보고.
    """
    L, gdp, exo, Wio = load(levels_csv, exog_hist_csv, io_csv)
    inds = list(L.columns); qs = list(L.index)
    tr_idx = qs.index(train_end)
    Ltr = L.iloc[:tr_idx + 1]
    val_q = [q for q in qs if "2020" <= q[:4] <= "2025"]
    Y = Ltr.diff(); sig = ar_sigma(Y.dropna()); Ysub = Y.dropna()
    Ym, X, groups = design(Ysub, exo.loc[Ysub.index], True, None)
    V = prior_var(groups, sig, len(inds), "io", Wio)
    B = fit_eqwise(Ym, X, V, sig, delta1=False)
    efull = exo[EXOG]
    Ld = L.diff()
    Gv = {q: Ld.loc[q].values.astype(float) for q in Ltr.index}
    Lv = {q: Ltr.loc[q].values.astype(float) for q in Ltr.index}
    allq = list(Ltr.index) + val_q
    for t in val_q:
        ti = allq.index(t); cols = []
        for l in range(1, P + 1):
            cols.append(Gv[allq[ti - l]])
        for s in range(0, Q + 1):
            cols.append(efull.loc[allq[ti - s]].values.astype(float))
        cols.append(np.array([1.0]))
        yhat = np.concatenate([np.atleast_1d(c) for c in cols]) @ B
        Lv[t] = Lv[allq[ti - 1]] + yhat; Gv[t] = yhat
    model_log = pd.DataFrame({t: Lv[t] for t in val_q}, index=inds).T
    gbar = (Ltr.iloc[-1] - Ltr.iloc[0]) / (len(Ltr) - 1)
    last = Ltr.iloc[-1]
    trend_log = pd.DataFrame({t: (last + gbar * (allq.index(t) - tr_idx)).values for t in val_q}, index=inds).T
    act_log = L.loc[val_q]

    def to_annual_level(dlog):
        lvl = np.exp(dlog / 100.0)
        yr = pd.Series([int(q[:4]) for q in dlog.index], index=dlog.index)
        return lvl.groupby(yr).sum()

    A_act = to_annual_level(act_log)

    def blend_rmse(w, years):
        A = to_annual_level(w * trend_log + (1 - w) * model_log)
        return float(np.sqrt(np.mean([((A.loc[y] - A_act.loc[y]) ** 2).mean() for y in years])))

    ws = np.round(np.arange(0, 1.001, 0.05), 2)
    val = [(float(w), blend_rmse(w, val_years)) for w in ws]
    wstar = min(val, key=lambda x: x[1])[0]
    tst = [(float(w), blend_rmse(w, test_years)) for w in ws]
    return {"ws": [w for w, _ in val], "val": [r for _, r in val], "test_curve": [r for _, r in tst],
            "wstar": float(wstar),
            "test": {"model": blend_rmse(0, test_years), "shrink": blend_rmse(wstar, test_years),
                     "trend": blend_rmse(1, test_years)}}