# -*- coding: utf-8 -*-
"""
Chapter 3 — BVARX (외생변수 결합)
 y_t = c + Σ_{l=1..p} B_l y_{t-l} + Σ_{s=0..q} Γ_s x_{t-s} + u_t
 내생 y: 36산업 성장률(미네소타 prior),  외생 x: 유가·교역·실질금리·실질환율·노동(느슨 prior)
 (1) 외생블록이 적합/예측을 높이나? — BVAR vs BVARX, 동일 test 조건부 1단계 RMSE
 (2) 외생충격의 산업별 동태승수 D_h = ∂y_{t+h}/∂x_t (20분기 누적, 1σ)
산출: ch3/figures/(fig1_bvar_vs_bvarx.png, fig2_multipliers.png, fig3_oil_response.png),
      ch3/data/(입력사본, ch3_results.json, exog_multipliers.csv)
"""
import os, json, shutil
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__)); CH3=os.path.dirname(HERE); REPO=os.path.dirname(CH3)
DATA=os.path.join(CH3,"data"); os.makedirs(DATA,exist_ok=True)
FIG=os.path.join(CH3,"figures"); os.makedirs(FIG,exist_ok=True)
NB,NR,NG="#1f4e79","#c0392b","#2e7d32"
P,Q,LAM,LAMX,WIN,TESTN = 4,2,0.1,1.0,30.0,20
EXO=["oil_g","trd_g","d_rrate","rfx_g","lab_g"]
EXOK=["유가","교역","실질금리","실질환율","노동"]

# --- 입력 복사(자기완결) ---
shutil.copy(os.path.join(REPO,"ch1","data","dlog_rates_sa.csv"),os.path.join(DATA,"dlog_rates_sa.csv"))
shutil.copy(os.path.join(REPO,"data","exog_quarterly_filled.csv"),os.path.join(DATA,"exog_quarterly_filled.csv"))

# --- 공통표본 ---
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)
ex=pd.read_csv(os.path.join(DATA,"exog_quarterly_filled.csv"),encoding="utf-8-sig")
ex.index=pd.PeriodIndex(ex["분기"],freq="Q"); ex=ex[EXO].astype(float)
common=g.index.intersection(ex.index).sort_values()
g=g.loc[common]; ex=ex.loc[common]
names=list(g.columns); n=len(names); ne=len(EXO)
gv=g.values; xv=ex.values; T=gv.shape[0]; test_start=T-TESTN
sig_x_full=xv.std(0)

def ar_sigma(m,p):
    s=np.zeros(m.shape[1])
    for i in range(m.shape[1]):
        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 design(gvv,xvv,p,q,use_x):
    T0=max(p,q); TT=len(gvv)-T0; blocks=[]
    for l in range(1,p+1): blocks.append(gvv[[t-l for t in range(T0,len(gvv))]])
    if use_x:
        for s in range(0,q+1): blocks.append(xvv[[t-s for t in range(T0,len(gvv))]])
    X=np.column_stack(blocks+[np.ones(TT)]); Y=gvv[T0:]; return Y,X

def penalty(sig_y,sig_x,p,q,use_x):
    d=[]
    for l in range(1,p+1):
        for j in range(n): d.append((l**2)*(sig_y[j]**2)/LAM**2)
    if use_x:
        for s in range(0,q+1):
            for m in range(ne): d.append(((s+1)**2)*(sig_x[m]**2)/LAMX**2)
    d.append(1e-6); return np.array(d)

def fit(gvv,xvv,p,q,use_x):
    Y,X=design(gvv,xvv,p,q,use_x); sy=ar_sigma(gvv,p); sx=xvv.std(0)
    d=penalty(sy,sx,p,q,use_x)
    B=np.linalg.solve(X.T@X+np.diag(d),X.T@Y); return B

def fc1(B,gvv,xvv,tgt,p,q,use_x):    # 조건부 1단계(외생 실제값 사용)
    x=[]
    for l in range(1,p+1): x.append(gvv[tgt-l])
    if use_x:
        for s in range(0,q+1): x.append(xvv[tgt-s])
    x=np.concatenate(x+[[1.0]]); return x@B

# --- 조건부 1단계 OOS: BVAR vs BVARX ---
def oos(use_x):
    err=[]
    for k in range(test_start,T):
        o=k-1
        if o<max(P,Q)+2: continue
        B=fit(gv[:k],xv[:k],P,Q,use_x)   # 학습은 target 직전까지(외생도 실제 관측)
        err.append(gv[k]-fc1(B,gv,xv,k,P,Q,use_x))
    E=np.array(err); return np.sqrt((E**2).mean()), np.sqrt((E**2).mean(0))
rm_bvar=oos(False); rm_bvarx=oos(True)
win=int((rm_bvarx[1]<rm_bvar[1]).sum())

# --- 전표본 BVARX 동태승수 ---
B=fit(gv,xv,P,Q,True)
A=[B[(l-1)*n:l*n,:].T for l in range(1,P+1)]
G=[B[n*P+s*ne:n*P+(s+1)*ne,:].T for s in range(Q+1)]   # (n x ne)
H=20; D=[]
for h in range(H+1):
    Dh=G[h].copy() if h<=Q else np.zeros((n,ne))
    for l in range(1,P+1):
        if h-l>=0: Dh=Dh+A[l-1]@D[h-l]
    D.append(Dh)
cum=np.sum(D,axis=0)*sig_x_full[None,:]     # 20분기 누적, 1σ  (n x ne)
Cdf=pd.DataFrame(cum,index=names,columns=EXOK)
Cdf.round(3).to_csv(os.path.join(DATA,"exog_multipliers.csv"),encoding="utf-8-sig")

def short(s): return (s.replace(" 제조업","").replace(" 및 ","·").replace("서비스업","서비스"))[:9]
print("="*58); print("챕터3 — BVAR vs BVARX  (조건부 1단계, test 2021Q1~2025Q4)"); print("="*58)
print(f"  BVAR  pooled RMSE = {rm_bvar[0]:.3f}")
print(f"  BVARX pooled RMSE = {rm_bvarx[0]:.3f}   (개선 {(1-rm_bvarx[0]/rm_bvar[0])*100:+.1f}%, 우세 {win}/{n} 산업)")
print("\n[유가 1σ↑ 20분기 누적반응 — 상위/하위]")
oil=Cdf["유가"].sort_values()
for nm,v in list(oil.items())[-3:][::-1]: print(f"  수혜 {v:+.2f}  {short(nm)}")
for nm,v in list(oil.items())[:3]: print(f"  피해 {v:+.2f}  {short(nm)}")

# --- 표본내 적합 R² (H3 검증) ---
def insample_r2(use_x):
    Y,X=design(gv,xv,P,Q,use_x); Bf=fit(gv,xv,P,Q,use_x); r=Y-X@Bf
    return float(np.mean(1-r.var(0)/Y.var(0)))
r2b=insample_r2(False); r2x=insample_r2(True)
print(f"\n표본내 R²: BVAR {r2b:.3f} -> BVARX {r2x:.3f}  (외생블록으로 +{(r2x-r2b)*100:.1f}%p, 논문 H3의 표본내 개선)")

# --- fig1: 표본내 R²(상승) vs 표본외 RMSE(불변) ---
fig,ax=plt.subplots(1,2,figsize=(11,4.4))
fig.suptitle("그림 1. 외생블록 — 표본내 적합은 오르나 표본외 예측은 개선 없음",fontweight="bold")
b0=ax[0].bar(["BVAR","BVARX"],[r2b,r2x],color=[NR,NB]); ax[0].bar_label(b0,fmt="%.3f")
ax[0].set_title("(a) 표본내 R² — 외생블록으로 상승 (H3)"); ax[0].set_ylabel("평균 R²")
b1=ax[1].bar(["BVAR","BVARX"],[rm_bvar[0],rm_bvarx[0]],color=[NR,NB]); ax[1].bar_label(b1,fmt="%.3f")
ax[1].set_title(f"(b) 표본외 1단계 RMSE — 대등/불변 (우세 {win}/{n})"); ax[1].set_ylabel("RMSE (%p)")
lo=min(rm_bvar[0],rm_bvarx[0]); ax[1].set_ylim(lo*0.97,max(rm_bvar[0],rm_bvarx[0])*1.02)
fig.tight_layout(rect=[0,0,1,.93]); fig.savefig(os.path.join(FIG,"fig1_bvar_vs_bvarx.png"),bbox_inches="tight"); plt.close(fig)

# --- fig2: 동태승수 히트맵 (산업 x 외생충격) ---
fig,ax=plt.subplots(figsize=(7,11))
v=np.percentile(np.abs(cum),95)
im=ax.imshow(cum,cmap="RdBu_r",norm=TwoSlopeNorm(0,-v,v),aspect="auto")
ax.set_xticks(range(ne)); ax.set_xticklabels(EXOK,fontsize=9)
ax.set_yticks(range(n)); ax.set_yticklabels([short(s) for s in names],fontsize=7)
ax.set_title("그림 2. 외생충격 1σ의 산업별 20분기 누적반응\n(붉음=수혜, 청=피해)",fontsize=11)
fig.colorbar(im,ax=ax,shrink=.7,pad=.02,label="%p")
fig.tight_layout(); fig.savefig(os.path.join(FIG,"fig2_multipliers.png"),bbox_inches="tight"); plt.close(fig)

# --- fig3: 유가충격 산업별 반응(비대칭) ---
fig,ax=plt.subplots(figsize=(8,9))
oil_s=Cdf["유가"].sort_values()
cols=[NR if x<0 else NG for x in oil_s.values]
ax.barh(range(n),oil_s.values,color=cols)
ax.set_yticks(range(n)); ax.set_yticklabels([short(s) for s in oil_s.index],fontsize=7)
ax.axvline(0,color="#333",lw=.8)
ax.set_title("그림 3. 유가 1σ↑ 충격의 산업별 반응 (비대칭)\n초록=수혜, 빨강=피해")
ax.set_xlabel("20분기 누적 %p")
fig.tight_layout(); fig.savefig(os.path.join(FIG,"fig3_oil_response.png"),bbox_inches="tight"); plt.close(fig)

json.dump({"rmse":{"BVAR":round(float(rm_bvar[0]),3),"BVARX":round(float(rm_bvarx[0]),3)},
           "r2":{"BVAR":round(r2b,3),"BVARX":round(r2x,3)},
           "improve_pct":round((1-rm_bvarx[0]/rm_bvar[0])*100,1),"win":f"{win}/{n}",
           "oil_top":[[short(k),round(float(v),2)] for k,v in oil.tail(3).items()],
           "oil_bottom":[[short(k),round(float(v),2)] for k,v in oil.head(3).items()]},
          open(os.path.join(DATA,"ch3_results.json"),"w",encoding="utf-8"),ensure_ascii=False,indent=2)
print("\n저장: ch3/figures/(fig1,fig2,fig3), ch3/data/(exog_multipliers.csv, ch3_results.json)")