# -*- coding: utf-8 -*-
"""
Chapter 4 — DSGE (BOKDPM 핵심 블록 이식): 시나리오별 거시·노동 경로 생성
 시나리오 충격(잠재성장·세계수요·유가·무역비용·인구) → 준구조 NK 모형(Fair-Taylor)
 → BVARX(챕터3) 외생 5종(oil_g, trd_g, d_rrate, rfx_g, lab_g) 산출.
 원 저장소 dsge_block1.py 이식 · 입력: kosis_생산연령인구.csv (KOSIS 장래인구추계)
산출: ch4/data/(dsge_exog_paths.csv, dsge_macro_paths.csv), ch4/figures/(fig1,fig2)
"""
import os, json, shutil
import numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

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__)); CH4=os.path.dirname(HERE); REPO=os.path.dirname(CH4)
DATA=os.path.join(CH4,"data"); os.makedirs(DATA,exist_ok=True)
FIG=os.path.join(CH4,"figures"); os.makedirs(FIG,exist_ok=True)
shutil.copy(os.path.join(REPO,"data","kosis_생산연령인구.csv"),os.path.join(DATA,"kosis_생산연령인구.csv"))

# --- BOKDPM calibration ---
b1,b2,b3,b4,b6=0.50,0.10,0.10,0.30,0.10
lam1,lam2,lam3=0.55,0.15,0.65
g1,g2,g3=0.50,1.50,0.50
a1,a2=0.80,0.30
om4,om7,om8=0.25,0.50,0.20
rrbar,pietar=1.5,2.0
H,BUF=40,24; N=H+BUF
SCN=["S1","S2","S3","S4","S5"]
SCN_KO={"S1":"기준","S2":"공급제약+충격","S3":"제약완화","S4":"기술도약","S5":"지정학"}

def labor_force_growth():
    pop=pd.read_csv(os.path.join(DATA,"kosis_생산연령인구.csv"),encoding="utf-8-sig").set_index("연도")
    col={"중위":"중위_천명","고위":"고위_천명","저위":"저위_천명"}; out={}
    for a,c in col.items():
        ann={y:np.log(pop.loc[y,c]/pop.loc[y-1,c])*100 for y in range(2026,2036)}
        out[a]=np.array([ann[2026+h//4]/4 for h in range(H)])
    return out

def scenario_inputs(s):
    z=np.zeros(N); dec=lambda k:np.clip(1-np.arange(N)/k,0,1)
    yus=z.copy(); oil=z.copy(); tcost=z.copy(); dshk=z.copy()
    if s=="S2": yus=-1.2*dec(12); oil=-8.0*dec(8); tcost=+1.5*dec(12); dshk=-1.0*dec(12)
    if s=="S4": yus=+0.8*dec(20); dshk=+0.6*np.ones(N)*dec(28)
    if s=="S5": oil=+15.0*dec(8); tcost=+2.0*dec(10); yus=-0.6*dec(10); dshk=-0.8*dec(10)
    return yus,oil,tcost,dshk

def solve_korea_block(yus,oil,dshk,iters=400):
    Y=np.zeros(N); PIE=np.full(N,pietar); RS=np.full(N,rrbar+pietar); S=np.zeros(N)
    for _ in range(iters):
        Y0,PIE0,RS0,S0=Y.copy(),PIE.copy(),RS.copy(),S.copy()
        for t in range(N):
            yl=Y[t-1] if t>0 else 0.0; yf=Y[t+1] if t+1<N else 0.0
            rsl=RS[t-1] if t>0 else rrbar+pietar
            rgap_l=(RS[t-1] if t>0 else rrbar+pietar)-pietar-rrbar
            Y[t]=(b1*yl+b2*yf-b3*rgap_l+b4*yus[t]-b6*oil[t]/10+dshk[t])
            pif=PIE[t+1] if t+1<N else pietar; pil=PIE[t-1] if t>0 else pietar
            PIE[t]=lam3*pietar+(1-lam3)*(lam1*pif+(1-lam1)*pil+lam2*4*yl)
            pie_e=PIE[t+1] if t+1<N else pietar
            RS[t]=g1*rsl+(1-g1)*(rrbar+pietar+g2*(pie_e-pietar)+g3*4*Y[t])
            sf=S[t+1] if t+1<N else 0.0; sl=S[t-1] if t>0 else 0.0
            rgap=RS[t]-pietar-rrbar
            S[t]=(1-0.1)*(om7*sf+(1-om7)*sl-om4*rgap)+om8*(-Y[t])
        if max(np.abs(Y-Y0).max(),np.abs(RS-RS0).max(),np.abs(S-S0).max())<1e-8: break
    RR=RS-PIE
    return Y[:H],PIE[:H],RS[:H],RR[:H],S[:H]

lfg=labor_force_growth()
lab_assump={"S1":"중위","S2":"저위","S3":"고위","S4":"중위","S5":"중위"}
partic=np.linspace(0.05,0.0,H)
quarters=[f"{2026+h//4}Q{h%4+1}" for h in range(H)]
rows,macro=[],[]
for s in SCN:
    yus,oil_sh,tcost,dshk=scenario_inputs(s)
    Y,PIE,RS,RR,S=solve_korea_block(yus,oil_sh,dshk)
    UNRgap=np.zeros(H)
    for t in range(H): UNRgap[t]=a1*(UNRgap[t-1] if t>0 else 0)-a2*Y[t]
    dUNR=np.diff(np.concatenate([[0],UNRgap]))
    oil_g=0.5+oil_sh[:H]/8*4
    trd_g=1.0+0.8*yus[:H]-0.6*tcost[:H]
    d_rrate=np.diff(np.concatenate([[RR[0]],RR]))
    rfx_g=np.diff(np.concatenate([[0],S]))
    lab_g=lfg[lab_assump[s]]+partic-dUNR
    for h in range(H):
        rows.append([s,quarters[h],round(oil_g[h],3),round(trd_g[h],3),round(d_rrate[h],3),round(rfx_g[h],3),round(lab_g[h],3)])
        macro.append([s,quarters[h],round(Y[h],3),round(PIE[h],3),round(RS[h],3),round(RR[h],3),round(S[h],3),round(UNRgap[h],3)])
ex=pd.DataFrame(rows,columns=["시나리오","분기","oil_g","trd_g","d_rrate","rfx_g","lab_g"])
mac=pd.DataFrame(macro,columns=["시나리오","분기","산출갭Y","인플레PIE","정책금리RS","실질금리RR","실질환율갭S","실업갭UNR"])
ex.to_csv(os.path.join(DATA,"dsge_exog_paths.csv"),index=False,encoding="utf-8-sig")
mac.to_csv(os.path.join(DATA,"dsge_macro_paths.csv"),index=False,encoding="utf-8-sig")

# --- 콘솔 ---
print("챕터4 — DSGE(BOKDPM 포트) 시나리오 경로 생성 완료")
print(f"{'시나리오':<16}{'2030 산출갭':>10}{'10년 고용Σ':>10}{'10년 실질금리Δ':>13}")
for s in SCN:
    m=mac[mac['시나리오']==s]; e=ex[ex['시나리오']==s]
    print(f"{s+' '+SCN_KO[s]:<16}{m['산출갭Y'].iloc[16]:>10.2f}{e['lab_g'].sum():>10.2f}{e['d_rrate'].sum():>13.2f}")

# --- fig1: 5 외생 경로 (BVARX 입력) ---
COLS=[("oil_g","유가 증가율"),("trd_g","교역량 증가율"),("d_rrate","실질금리 변화"),("rfx_g","실질환율 변화"),("lab_g","고용 증가율")]
CLR={"S1":"#333","S2":"#c0392b","S3":"#1f4e79","S4":"#2e7d32","S5":"#b8892f"}
t=[pd.Period(q,"Q").to_timestamp(how="end") for q in quarters]
fig,axes=plt.subplots(2,3,figsize=(15,7.5))
fig.suptitle("그림 1. DSGE가 생성한 시나리오별 외생 경로 (BVARX 입력 5종, 2026–2035)",fontsize=14,fontweight="bold")
for ax,(c,lab) in zip(axes.ravel(),COLS):
    for s in SCN:
        d=ex[ex['시나리오']==s]
        ax.plot(t,d[c].values,color=CLR[s],lw=1.6,label=f"{s} {SCN_KO[s]}")
    ax.axhline(0,color="grey",lw=.6); ax.set_title(lab); ax.grid(alpha=.25); ax.margins(x=.02)
axes.ravel()[-1].axis("off")
axes.ravel()[0].legend(loc="center left",bbox_to_anchor=(3.25,-.4),fontsize=10,frameon=False)
fig.tight_layout(rect=[0,0,1,.95]); fig.savefig(os.path.join(FIG,"fig1_exog_paths.png"),bbox_inches="tight"); plt.close(fig)

# --- fig2: 거시 산출갭 + 고용 ---
fig,ax=plt.subplots(1,2,figsize=(13,4.6))
fig.suptitle("그림 2. DSGE 거시 블록 — 시나리오별 산출갭과 고용",fontweight="bold")
for s in SCN:
    m=mac[mac['시나리오']==s]; e=ex[ex['시나리오']==s]
    ax[0].plot(t,m['산출갭Y'].values,color=CLR[s],lw=1.6,label=f"{s} {SCN_KO[s]}")
    ax[1].plot(t,e['lab_g'].cumsum().values,color=CLR[s],lw=1.6)
ax[0].axhline(0,color="grey",lw=.6); ax[0].set_title("(a) 산출갭 (%)"); ax[0].legend(fontsize=8); ax[0].grid(alpha=.25)
ax[1].set_title("(b) 누적 고용증가 (%) — 인구 가정 차이"); ax[1].grid(alpha=.25)
fig.tight_layout(rect=[0,0,1,.93]); fig.savefig(os.path.join(FIG,"fig2_macro.png"),bbox_inches="tight"); plt.close(fig)

json.dump({s:{"lab_sum":round(float(ex[ex['시나리오']==s]['lab_g'].sum()),2),
              "Y2030":round(float(mac[mac['시나리오']==s]['산출갭Y'].iloc[16]),2)} for s in SCN},
          open(os.path.join(DATA,"ch4_results.json"),"w",encoding="utf-8"),ensure_ascii=False,indent=2)
print("저장: ch4/data/(dsge_exog_paths.csv, dsge_macro_paths.csv), ch4/figures/(fig1,fig2)")