# -*- coding: utf-8 -*-
"""
Chapter 5 — 시나리오 시뮬레이션 (S1–S5)
 챕터 3의 BVARX에 챕터 4의 DSGE 외생 경로를 '조건화'해 36개 산업 10년 경로를 낸다.
 = 외생 경로(유가·교역·금리·환율·고용)를 넣고 BVARX를 앞으로 굴려 산업별 성장 → 수준.
산출: ch5/figures/(fig1_gdp_fan.png, fig2_industry.png), ch5/data/(scenario_industry_2035.csv, ch5_results.json)
"""
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__)); CH5=os.path.dirname(HERE); REPO=os.path.dirname(CH5)
DATA=os.path.join(CH5,"data"); os.makedirs(DATA,exist_ok=True)
FIG=os.path.join(CH5,"figures"); os.makedirs(FIG,exist_ok=True)
P,Q,LAM,LAMX,WIN=4,2,0.1,1.0,30.0
EXO=["oil_g","trd_g","d_rrate","rfx_g","lab_g"]
SCN=["S1","S2","S3","S4","S5"]; SCN_KO={"S1":"기준","S2":"공급제약+충격","S3":"제약완화","S4":"기술도약","S5":"지정학"}
CLR={"S1":"#333","S2":"#c0392b","S3":"#1f4e79","S4":"#2e7d32","S5":"#b8892f"}

# 입력 복사
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"))
shutil.copy(os.path.join(REPO,"ch4","data","dsge_exog_paths.csv"),os.path.join(DATA,"dsge_exog_paths.csv"))
shutil.copy(os.path.join(REPO,"..","analysis","prepared","levels_sa.csv"),os.path.join(DATA,"levels_sa.csv"))

# 데이터·공통표본 (챕터3과 동일)
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

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):
    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))]])
    for s in range(0,q+1): blocks.append(xvv[[t-s for t in range(T0,len(gvv))]])
    return gvv[T0:], np.column_stack(blocks+[np.ones(TT)])
# BVARX 추정 (전표본)
sy=ar_sigma(gv,P); sx=xv.std(0)
d=[]
for l in range(1,P+1):
    for j in range(n): d.append((l**2)*(sy[j]**2)/LAM**2)
for s in range(0,Q+1):
    for m in range(ne): d.append(((s+1)**2)*(sx[m]**2)/LAMX**2)
d.append(1e-6); d=np.array(d)
Y,X=design(gv,xv,P,Q); B=np.linalg.solve(X.T@X+np.diag(d),X.T@Y)

# DSGE 시나리오 외생 경로
dsge=pd.read_csv(os.path.join(DATA,"dsge_exog_paths.csv"),encoding="utf-8-sig")
Hf=dsge[dsge["시나리오"]=="S1"].shape[0]      # 40분기

def scen_forecast(scn):
    xs=dsge[dsge["시나리오"]==scn][EXO].values          # (Hf, ne)
    exog_full=np.vstack([xv, xs]); start=len(xv)
    buf=[r.copy() for r in gv[-P:]]; out=[]
    for h in range(Hf):
        ti=start+h; row=[]
        for l in range(1,P+1): row.append(buf[-l])
        for s in range(0,Q+1): row.append(exog_full[ti-s])
        y=np.concatenate(row+[[1.0]])@B; buf.append(y); out.append(y)
    return np.array(out)                                   # (Hf, n) 성장률

# 가중치(2025Q4 실질수준)
lev=pd.read_csv(os.path.join(DATA,"levels_sa.csv"),encoding="utf-8-sig",index_col=0)
w=lev[names].iloc[-1].values; w=w/w.sum()

gdp_idx={}; ind2035={}
for s in SCN:
    gr=scen_forecast(s)                                    # (Hf,n) %
    agg=(gr*w[None,:]).sum(1)                              # 가중 총량 성장 %
    gdp_idx[s]=100*np.exp(np.cumsum(agg/100))
    ind2035[s]=100*np.exp(np.cumsum(gr/100,0))[-1]         # 산업 2035 지수(2025=100)
ind_df=pd.DataFrame(ind2035,index=names)[SCN]
ind_df.round(1).to_csv(os.path.join(DATA,"scenario_industry_2035.csv"),encoding="utf-8-sig")

def short(s): return (s.replace(" 제조업","").replace(" 및 ","·").replace("서비스업","서비스"))[:9]
print("챕터5 — 시나리오 시뮬레이션 (BVARX×DSGE 조건화)")
print(f"{'시나리오':<16}{'2035 GDP지수(2025=100)':>22}")
for s in SCN: print(f"{s+' '+SCN_KO[s]:<16}{gdp_idx[s][-1]:>22.1f}")
labor=(ind_df["S3"]-ind_df["S2"]).sort_values()          # 제약완화-공급제약 = 노동효과
print("\n[노동효과(S3−S2) 큰 산업]")
for nm,v in labor.tail(4)[::-1].items(): print(f"  +{v:.1f}  {short(nm)}")

t=[pd.Period(q,"Q").to_timestamp(how="end") for q in dsge[dsge['시나리오']=='S1']['분기']]
# fig1: GDP 지수 시나리오 팬
fig,ax=plt.subplots(figsize=(9,5))
for s in SCN: ax.plot(t,gdp_idx[s],color=CLR[s],lw=2,label=f"{s} {SCN_KO[s]} ({gdp_idx[s][-1]:.0f})")
ax.set_title("그림 1. 시나리오별 총량 실질 GDP 경로 (2025Q4=100)",fontweight="bold")
ax.set_ylabel("지수 (2025=100)"); ax.legend(fontsize=9); ax.grid(alpha=.3); ax.margins(x=.02)
fig.tight_layout(); fig.savefig(os.path.join(FIG,"fig1_gdp_fan.png"),bbox_inches="tight"); plt.close(fig)

# fig2: 산업별 노동효과(S3-S2)
fig,ax=plt.subplots(figsize=(8,9)); ls=labor
cols=["#c0392b" if x<0 else "#1f4e79" for x in ls.values]
ax.barh(range(n),ls.values,color=cols); ax.axvline(0,color="#333",lw=.8)
ax.set_yticks(range(n)); ax.set_yticklabels([short(x) for x in ls.index],fontsize=7)
ax.set_title("그림 2. 노동공급 효과 = 제약완화(S3) - 공급제약(S2)\n2035 산업 지수 차이 (클수록 인구에 민감)")
ax.set_xlabel("2035 지수 차이 (%p)")
fig.tight_layout(); fig.savefig(os.path.join(FIG,"fig2_industry.png"),bbox_inches="tight"); plt.close(fig)

json.dump({"gdp2035":{s:round(float(gdp_idx[s][-1]),1) for s in SCN},
           "labor_top":[[short(k),round(float(v),1)] for k,v in labor.tail(4)[::-1].items()]},
          open(os.path.join(DATA,"ch5_results.json"),"w",encoding="utf-8"),ensure_ascii=False,indent=2)
print("\n저장: ch5/figures/(fig1,fig2), ch5/data/(scenario_industry_2035.csv, ch5_results.json)")