# -*- coding: utf-8 -*-
"""
Chapter 2 — 미네소타 prior vs 산업연관표(I-O) prior BVAR 비교
 · 추정: 방정식별 Bayesian ridge (I-O prior가 계수(i←j)를 ω_ij로 차등 → 켤레구조 불가)
   계수(l,j) 릿지 벌점 d:  미네소타 = l²σ_j²/λ²,   I-O(교차 j≠i) = l²σ_j²/(λ²·ω_ij)
 · ω: io_prior_omega.csv (오프대각 평균 1로 재정규화 → 총수축 동일, 배분만 차등)
 · 검증: 동일 test 구간(2021Q1–2025Q4) 확장창 1단계·1년(4Q) 예측 RMSE
산출: ch2/figures/(fig2_rmse_compare.png, fig3_coef_structure.png), ch2/data/ch2_results.json
"""
import os, json
import numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm

plt.rcParams["font.family"]="Malgun Gothic"; plt.rcParams["axes.unicode_minus"]=False
plt.rcParams["savefig.dpi"]=130
HERE=os.path.dirname(os.path.abspath(__file__)); CH2=os.path.dirname(HERE)
DATA=os.path.join(CH2,"data"); FIG=os.path.join(CH2,"figures")
P, LAM, WIN, TESTN = 4, 0.1, 30.0, 20
NB, NR, NG = "#1f4e79", "#c0392b", "#2e7d32"

# --- 데이터 ---
rate=pd.read_csv(os.path.join(DATA,"dlog_rates_sa.csv"),encoding="utf-8-sig",index_col=0)
rate.index=pd.PeriodIndex(rate.index,freq="Q")
g=rate.drop(columns=["총량GDP"]).dropna().clip(-WIN,WIN)
names=list(g.columns); n=len(names); gv=g.values; T=gv.shape[0]
test_start=T-TESTN
Om=pd.read_csv(os.path.join(DATA,"io_prior_omega.csv"),encoding="utf-8-sig",index_col=0).values
off=~np.eye(n,dtype=bool)
Om[off]=Om[off]/Om[off].mean()          # 오프대각 평균 1로 재정규화(공정)
np.fill_diagonal(Om,1.0)

def build_XY(m,p):
    Y=m[p:]; TT=Y.shape[0]; X=np.ones((TT,n*p+1))
    for l in range(1,p+1): X[:,(l-1)*n:l*n]=m[p-l:-l]
    return Y,X
def ar_sigma(m,p):
    s=np.zeros(n)
    for i in range(n):
        y=m[:,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 penalty_base(sig):
    """릿지 벌점 (l,j) = l²σ_j²/λ²  (미네소타, 방정식 공통). const≈0."""
    d=np.zeros(n*P+1)
    for l in range(1,P+1):
        for j in range(n): d[(l-1)*n+j]=(l**2)*(sig[j]**2)/(LAM**2)
    d[-1]=1e-6
    return d

def fit(m,p,mode,sig):
    Y,X=build_XY(m,p); XtX=X.T@X; XtY=X.T@Y
    d0=penalty_base(sig)
    if mode=="MIN":
        B=np.linalg.solve(XtX+np.diag(d0),XtY)         # 방정식 공통
    else:  # I-O: 방정식별 (교차항 벌점 /= ω_ij)
        B=np.zeros((n*P+1,n))
        for i in range(n):
            d=d0.copy()
            for l in range(1,p+1):
                for j in range(n):
                    if j!=i: d[(l-1)*n+j]/=Om[i,j]      # 연계 강하면(ω>1) 벌점↓
            B[:,i]=np.linalg.solve(XtX+np.diag(d),XtY[:,i])
    return B

def forecast(B,m,p,h):
    buf=[r.copy() for r in m[-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 evaluate(mode):
    res={1:[],4:[]}
    for k in range(test_start,T):
        for h in (1,4):
            o=k-h
            if o<P+4: continue
            tr=gv[:o+1]; sig=ar_sigma(tr,P); B=fit(tr,P,mode,sig)
            res[h].append((gv[k], forecast(B,tr,P,h)))
    return res
EV={m:evaluate(m) for m in ["MIN","IO"]}
def rmse(mode,h):
    a=np.array([x[0] for x in EV[mode][h]]); f=np.array([x[1] for x in EV[mode][h]])
    return np.sqrt(((a-f)**2).mean()), np.sqrt(((a-f)**2).mean(0))  # pooled, per-industry
RM={m:{h:rmse(m,h) for h in (1,4)} for m in ["MIN","IO"]}

# --- 전표본 계수 구조 (fig3) ---
sig_full=ar_sigma(gv,P)
Bmin=fit(gv,P,"MIN",sig_full); Bio=fit(gv,P,"IO",sig_full)
A1min=Bmin[0:n,:].T; A1io=Bio[0:n,:].T           # [반응 i, 충격 j]

# --- 콘솔 ---
print("="*60)
print("챕터2 — 미네소타 vs I-O prior  (동일 test, 방정식별 릿지)")
print("="*60)
print(f"{'':10}{'1단계 pooled':>14}{'1년 pooled':>12}")
for m,lb in [("MIN","미네소타"),("IO","I-O prior")]:
    print(f"{lb:10}{RM[m][1][0]:>14.3f}{RM[m][4][0]:>12.3f}")
win1=int((RM['IO'][1][1]<RM['MIN'][1][1]).sum()); win4=int((RM['IO'][4][1]<RM['MIN'][4][1]).sum())
print(f"\nI-O가 미네소타보다 나은 산업: 1단계 {win1}/{n}, 1년 {win4}/{n}")
print(f"교차 |A1| 평균: 미네소타 {np.abs(A1min[off]).mean():.4f}  I-O {np.abs(A1io[off]).mean():.4f}")

# --- fig2: test RMSE 비교 (표와 동일수치) ---
fig,ax=plt.subplots(1,2,figsize=(13,4.6))
fig.suptitle("그림 2. 미네소타 vs I-O prior — 동일 test 구간 예측 RMSE",fontweight="bold")
x=np.arange(2); w=.38
for hi,h in enumerate((1,4)):
    vals=[RM["MIN"][h][0],RM["IO"][h][0]]
    b=ax[hi].bar(["미네소타","I-O prior"],vals,color=[NR,NB])
    ax[hi].bar_label(b,fmt="%.3f")
    ax[hi].set_title(f"({'a' if h==1 else 'b'}) {'1단계' if h==1 else '1년앞(4Q)'} pooled RMSE")
    ax[hi].set_ylabel("RMSE (%p)")
    lo=min(vals); ax[hi].set_ylim(lo*0.985,max(vals)*1.01)
fig.tight_layout(rect=[0,0,1,.93]); fig.savefig(os.path.join(FIG,"fig2_rmse_compare.png"),bbox_inches="tight"); plt.close(fig)

# --- fig3: 계수 구조 (미네소타 vs I-O |A1|) ---
def short(s): return (s.replace(" 제조업","").replace(" 및 ","·").replace("서비스업","서비스"))[:9]
lab=[short(s) for s in names]
v=np.percentile(np.abs(np.r_[A1min[off],A1io[off]]),98)
fig,ax=plt.subplots(1,2,figsize=(15,7.2))
fig.suptitle("그림 3. 1차 계수 |A₁| 구조 — I-O prior는 연계 약한 쌍을 더 0으로 (열 j→행 i)",fontweight="bold")
for a_,Mx,tt in [(ax[0],np.abs(A1min),"(a) 미네소타"),(ax[1],np.abs(A1io),"(b) I-O prior")]:
    im=a_.imshow(Mx,cmap="YlOrRd",vmin=0,vmax=v,aspect="auto"); a_.set_title(tt)
    a_.set_xticks(range(n)); a_.set_xticklabels(lab,rotation=90,fontsize=5)
    a_.set_yticks(range(n)); a_.set_yticklabels(lab,fontsize=5)
    fig.colorbar(im,ax=a_,shrink=.8,pad=.02)
fig.tight_layout(rect=[0,0,1,.95]); fig.savefig(os.path.join(FIG,"fig3_coef_structure.png"),bbox_inches="tight"); plt.close(fig)

# --- 결과 저장 ---
json.dump({"rmse":{m:{str(h):round(float(RM[m][h][0]),3) for h in (1,4)} for m in ["MIN","IO"]},
           "win_io":{"h1":f"{win1}/{n}","h4":f"{win4}/{n}"},
           "absA1_cross":{"MIN":round(float(np.abs(A1min[off]).mean()),4),"IO":round(float(np.abs(A1io[off]).mean()),4)}},
          open(os.path.join(DATA,"ch2_results.json"),"w",encoding="utf-8"),ensure_ascii=False,indent=2)
print("\n저장: ch2/figures/(fig2,fig3), ch2/data/ch2_results.json")