# -*- coding: utf-8 -*-
"""
방법론 템플릿 분석: AR → VAR → BVAR, 동일 test 구간, 1단계·1년(4Q) 예측.
대상: 컴퓨터·전자(A) ~ 전기장비(B)  (성장률 상관 0.636)
+ 상관 기반 바텀업 군집 설계.
산출: report/methodology/*.png, report/methodology/results.json
"""
import os, json
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
from scipy.cluster.hierarchy import linkage, dendrogram
from scipy.spatial.distance import squareform

plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams["savefig.dpi"] = 120
plt.rcParams["axes.grid"] = True
plt.rcParams["grid.alpha"] = 0.3

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)   # ch1/
PREP = os.path.join(ROOT, "data")
OUT = os.path.join(ROOT, "figures"); os.makedirs(OUT, exist_ok=True)
A, B = "컴퓨터, 전자 및 광학기기 제조업", "전기장비 제조업"
LA, LB = "컴퓨터·전자", "전기장비"
P, LAM, WIN, TESTN = 4, 0.1, 30.0, 20     # 시차, 수축, 윈저, 테스트분기수
NB = "#1f4e79"; NR = "#c0392b"; NG = "#2e7d32"

rate = pd.read_csv(os.path.join(PREP, "dlog_rates_sa.csv"), encoding="utf-8-sig", index_col=0)
rate.index = pd.PeriodIndex(rate.index, freq="Q")
allg = rate.drop(columns=["총량GDP"]).dropna().clip(-WIN, WIN)
g = allg[[A, B]]
idx = g.index
m = g.values                      # (T,2)
T = m.shape[0]; n = 2
test_start = T - TESTN            # 테스트 타깃 시작 인덱스
corrAB = float(np.corrcoef(m[:, 0], m[:, 1])[0, 1])

# ---------- 모형별 h-step 예측 함수 (train으로 추정, dynamic iterate) ----------
def build_XY(mt, p):
    Y = mt[p:]; TT = Y.shape[0]; nn = mt.shape[1]
    X = np.ones((TT, nn*p+1))
    for l in range(1, p+1):
        X[:, (l-1)*nn:l*nn] = mt[p-l:-l]
    X[:, -1] = 1.0
    return Y, X

def ar_sigma(mt, p):
    s = np.zeros(mt.shape[1])
    for i in range(mt.shape[1]):
        y = mt[:, i]; Y = y[p:]
        Z = np.column_stack([y[p-l:len(y)-l] for l in range(1, p+1)] + [np.ones(len(Y))])
        b, *_ = np.linalg.lstsq(Z, Y, rcond=None); s[i] = (Y - Z@b).std(ddof=Z.shape[1])
    return s

def iterate(B, mt, p, h):
    buf = [r.copy() for r in mt[-p:]]
    for _ in range(h):
        x = np.ones(n*p+1)
        for l in range(1, p+1):
            x[(l-1)*n:l*n] = buf[-l]
        buf.append(x @ B)
    return buf[-1]

def fit_AR(mt, p, h):                      # 개별 AR(4), 각 계열 독립
    out = np.zeros(n)
    for i in range(n):
        y = mt[:, i]; Y = y[p:]
        Z = np.column_stack([y[p-l:len(y)-l] for l in range(1, p+1)] + [np.ones(len(Y))])
        b, *_ = np.linalg.lstsq(Z, Y, rcond=None)
        buf = list(y[-p:])
        for _ in range(h):
            xz = [buf[-l] for l in range(1, p+1)] + [1.0]; buf.append(float(np.dot(b, xz)))
        out[i] = buf[-1]
    return out

def fit_VAR(mt, p, h):                      # OLS VAR
    Y, X = build_XY(mt, p); B = np.linalg.lstsq(X, Y, rcond=None)[0]
    return iterate(B, mt, p, h), B

def fit_BVAR(mt, p, h, lam=LAM):            # 미네소타 BVAR
    Y, X = build_XY(mt, p); sig = ar_sigma(mt, p)
    om = np.zeros(n*p+1)
    for l in range(1, p+1):
        for j in range(n): om[(l-1)*n+j] = lam**2/(l**2*sig[j]**2)
    om[-1] = 1e6
    Ob = np.linalg.inv(np.diag(1/om) + X.T@X); B = Ob @ (X.T@Y)
    return iterate(B, mt, p, h), B

# ---------- 재귀(확장창) 테스트 평가: 1단계 & 1년(4Q) ----------
def evaluate(model):
    res = {1: {"act": [], "fc": [], "q": []}, 4: {"act": [], "fc": [], "q": []}}
    for k in range(test_start, T):                   # 타깃 인덱스
        for h in (1, 4):
            o = k - h                                # 예측 원점
            if o < P + 4:
                continue
            train = m[:o+1]
            if model == "AR":
                fc = fit_AR(train, P, h)
            elif model == "VAR":
                fc, _ = fit_VAR(train, P, h)
            else:
                fc, _ = fit_BVAR(train, P, h)
            res[h]["act"].append(m[k]); res[h]["fc"].append(fc); res[h]["q"].append(idx[k])
    return res

EV = {mo: evaluate(mo) for mo in ["AR", "VAR", "BVAR"]}

def rmse(model, h, i):
    a = np.array(EV[model][h]["act"])[:, i]; f = np.array(EV[model][h]["fc"])[:, i]
    return float(np.sqrt(np.mean((a-f)**2)))

RM = {mo: {h: [rmse(mo, h, i) for i in range(n)] for h in (1, 4)} for mo in EV}

# 전표본 VAR-OLS vs BVAR 계수(비교용)
_, Bvar = fit_VAR(m, P, 1); _, Bbvar = fit_BVAR(m, P, 1)

# ================= 그림 1: 데이터 + 상관 =================
tt = idx.to_timestamp(how="end")
fig, ax = plt.subplots(1, 2, figsize=(13, 4.3))
ax[0].axhline(0, color="grey", lw=.6)
ax[0].plot(tt, m[:, 0], color=NB, lw=.8, label=LA)
ax[0].plot(tt, m[:, 1], color=NR, lw=.8, alpha=.8, label=LB)
ax[0].axvspan(idx[test_start].to_timestamp(how="end"), tt[-1], color="orange", alpha=.12)
ax[0].text(idx[test_start].to_timestamp(how="end"), ax[0].get_ylim()[1]*.9, " Test", color="#b26a00", fontsize=9)
ax[0].set_title(f"(a) 두 산업 분기 성장률 (오렌지=Test 최근 {TESTN}분기)"); ax[0].set_ylabel("%"); ax[0].legend(fontsize=8)
ax[1].scatter(m[:, 0], m[:, 1], s=14, color=NB, alpha=.5, edgecolor="none")
ax[1].set_title(f"(b) 두 산업 성장률 산점 — 상관 {corrAB:.2f}"); ax[1].set_xlabel(LA+" %"); ax[1].set_ylabel(LB+" %")
fig.suptitle("그림 1. 데이터: 컴퓨터·전자 ~ 전기장비 (상관 높은 쌍)", fontweight="bold")
fig.tight_layout(rect=[0, 0, 1, 0.94]); fig.savefig(os.path.join(OUT, "fig1_data.png"), bbox_inches="tight"); plt.close(fig)

# ================= 그림 2/3: AR, VAR 예측 추종 =================
def plot_track(model, fname, title):
    fig, ax = plt.subplots(1, 2, figsize=(13, 4.3))
    for i, (lab) in enumerate([LA, LB]):
        q1 = [q.to_timestamp(how="end") for q in EV[model][1]["q"]]
        q4 = [q.to_timestamp(how="end") for q in EV[model][4]["q"]]
        a1 = np.array(EV[model][1]["act"])[:, i]
        f1 = np.array(EV[model][1]["fc"])[:, i]
        f4 = np.array(EV[model][4]["fc"])[:, i]
        ax[i].axhline(0, color="grey", lw=.6)
        ax[i].plot(q1, a1, "o-", ms=3, color="#333", label="실측")
        ax[i].plot(q1, f1, "s-", ms=3, color=NB, label=f"1단계 (RMSE {RM[model][1][i]:.2f})")
        ax[i].plot(q4, f4, "^--", ms=3, color=NR, label=f"1년앞 다이내믹 (RMSE {RM[model][4][i]:.2f})")
        ax[i].set_title(f"{lab}"); ax[i].set_ylabel("%"); ax[i].legend(fontsize=8)
    fig.suptitle(title, fontweight="bold")
    fig.tight_layout(rect=[0, 0, 1, 0.93]); fig.savefig(os.path.join(OUT, fname), bbox_inches="tight"); plt.close(fig)

plot_track("AR", "fig2_ar.png", "그림 2. AR(4) — 테스트 구간 1단계·1년앞 예측 추종")
plot_track("VAR", "fig3_var.png", "그림 3. VAR(4) — 테스트 구간 1단계·1년앞 예측 추종")

# ================= 그림 4: VAR vs BVAR 계수 =================
fig, ax = plt.subplots(1, 2, figsize=(12, 4.6))
vv = Bvar[:-1].ravel(); bb = Bbvar[:-1].ravel()
lim = max(abs(vv).max(), abs(bb).max())*1.1
ax[0].plot([-lim, lim], [-lim, lim], color="grey", ls="--", lw=1)
ax[0].axhline(0, color="#ccc", lw=.6); ax[0].axvline(0, color="#ccc", lw=.6)
ax[0].scatter(vv, bb, s=30, color=NB, alpha=.7)
ax[0].set_xlabel("VAR (OLS) 계수"); ax[0].set_ylabel("BVAR (미네소타) 계수")
ax[0].set_title("(a) 2변수에선 VAR과 BVAR 거의 동일\n(BVAR가 0쪽으로 살짝 수축)")
# 축약 요약: |계수| 평균
ax[1].bar(["VAR (OLS)", "BVAR"], [np.abs(vv).mean(), np.abs(bb).mean()], color=[NR, NB])
ax[1].set_ylabel("평균 |동학계수|")
ax[1].set_title(f"(b) 평균 계수 크기\nBVAR가 {(1-np.abs(bb).mean()/np.abs(vv).mean())*100:.0f}% 더 작음(수축)")
fig.suptitle("그림 4. VAR와 BVAR의 차이 = '수축(prior)' — 변수 적을 땐 미미", fontweight="bold")
fig.tight_layout(rect=[0, 0, 1, 0.92]); fig.savefig(os.path.join(OUT, "fig4_var_vs_bvar.png"), bbox_inches="tight"); plt.close(fig)

# ================= 그림 5: RMSE 비교 =================
fig, ax = plt.subplots(1, 2, figsize=(13, 4.3))
models = ["AR", "VAR", "BVAR"]; x = np.arange(len(models)); w = .35
for hi, h in enumerate((1, 4)):
    avg = [np.mean(RM[mo][h]) for mo in models]
    ax[hi].bar(x, avg, color=[NR, NB, NG])
    ax[hi].bar_label(ax[hi].containers[0], fmt="%.2f")
    ax[hi].set_xticks(x); ax[hi].set_xticklabels(models)
    ax[hi].set_title(f"({'a' if h==1 else 'b'}) {'1단계' if h==1 else '1년앞(4Q)'} 예측 평균 RMSE")
    ax[hi].set_ylabel("RMSE (%p)")
fig.suptitle("그림 5. 테스트 구간 예측정확도 — AR vs VAR vs BVAR", fontweight="bold")
fig.tight_layout(rect=[0, 0, 1, 0.93]); fig.savefig(os.path.join(OUT, "fig5_rmse.png"), bbox_inches="tight"); plt.close(fig)

# ================= 그림 6: 상관 군집(바텀업 설계) =================
C = allg.corr().values
names = list(allg.columns)
short = [s.replace(" 제조업", "").replace(" 및 ", "·").replace("서비스업", "서비스")[:10] for s in names]
D = 1 - C; np.fill_diagonal(D, 0); D = (D + D.T)/2
Z = linkage(squareform(D, checks=False), method="average")
fig, ax = plt.subplots(1, 2, figsize=(15, 7), gridspec_kw={"width_ratios": [1.1, 1]})
dn = dendrogram(Z, labels=short, ax=ax[0], orientation="left", color_threshold=0.72, leaf_font_size=7)
ax[0].set_title("(a) 성장률 상관 기반 위계군집 (바텀업 병합)")
ax[0].set_xlabel("거리 = 1 - 상관")
order = dn["leaves"]
Cn = C[np.ix_(order, order)]
im = ax[1].imshow(Cn, cmap="RdBu_r", norm=TwoSlopeNorm(0, -1, 1), aspect="auto")
ax[1].set_xticks(range(len(order))); ax[1].set_xticklabels([short[i] for i in order], rotation=90, fontsize=5)
ax[1].set_yticks(range(len(order))); ax[1].set_yticklabels([short[i] for i in order], fontsize=5)
ax[1].set_title("(b) 군집순 상관행렬 (블록=함께 묶을 후보)")
fig.colorbar(im, ax=ax[1], shrink=.8, pad=.02)
fig.suptitle("그림 6. 바텀업 설계: 상관 높은 산업을 블록으로 묶어 최종 BVAR로", fontweight="bold")
fig.tight_layout(rect=[0, 0, 1, 0.95]); fig.savefig(os.path.join(OUT, "fig6_cluster.png"), bbox_inches="tight"); plt.close(fig)

# ---------- 결과 JSON ----------
results = {
    "pair": [LA, LB], "corr": round(corrAB, 3), "p": P, "lam": LAM, "test_n": TESTN,
    "test_period": [str(idx[test_start]), str(idx[-1])],
    "train_end": str(idx[test_start-1]),
    "rmse": {mo: {str(h): [round(v, 3) for v in RM[mo][h]] for h in (1, 4)} for mo in models},
    "rmse_avg": {mo: {str(h): round(float(np.mean(RM[mo][h])), 3) for h in (1, 4)} for mo in models},
    "coef_abs_var": round(float(np.abs(Bvar[:-1]).mean()), 4),
    "coef_abs_bvar": round(float(np.abs(Bbvar[:-1]).mean()), 4),
    "top_pairs": None,
}
json.dump(results, open(os.path.join(OUT, "results.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=2)

print("=== 테스트 RMSE (낮을수록 좋음) ===")
print(f"{'모형':6}{'1단계('+LA[:6]+')':>16}{'1단계('+LB[:4]+')':>14}{'1년('+LA[:6]+')':>14}{'1년('+LB[:4]+')':>12}")
for mo in models:
    print(f"{mo:6}{RM[mo][1][0]:>16.2f}{RM[mo][1][1]:>14.2f}{RM[mo][4][0]:>14.2f}{RM[mo][4][1]:>12.2f}")
print(f"\n평균 |동학계수|: VAR={np.abs(Bvar[:-1]).mean():.3f}  BVAR={np.abs(Bbvar[:-1]).mean():.3f}")
print("\n저장: figures/ (fig1~6.png, results.json)")