"""
주식과 채권의 상관관계가 물가 상승률에 따라 어떻게 달라지는지 미국 월별 데이터로 확인한다

사용법:
    pip install pandas numpy matplotlib
    python stock_bond_inflation.py
    python stock_bond_inflation.py --synthetic     # 난수 데이터로 동작 확인

데이터: Robert Shiller 교수가 공개하는 미국 월별 데이터(S&P500 지수, 배당, 10년물 국채 금리, 소비자물가지수)를
        datasets/s-and-p-500 저장소가 CSV로 정리해 둔 사본. asset_allocation_basics.py 와 같은 파일이다.
  - 1953년 이전 금리는 연간 값을 달마다 고르게 나눠 이은 것이라 월별 채권 수익률로 쓸 수 없다. 그래서 1953년 2월부터 쓴다.
  - 배당 열이 2023년 6월에서 끝나므로 2023년 6월까지 쓴다.
  - 주식: 지수 가격 변화 + 배당/12. 채권: 지난달 금리로 발행한 10년 만기 국채를 이번 달 금리로 다시 평가한 가격 변화 + 한 달치 이자.
  - 물가 구간: 직전 달까지 12개월 소비자물가 상승률(그 달 수익률을 보기 전에 알 수 있는 값)로 나눈다.
결과: output/stock_bond_inflation.txt 와 차트 PNG
"""
import argparse
import os

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "output")
URL = "https://raw.githubusercontent.com/datasets/s-and-p-500/main/data/data.csv"
START, END = "1953-01", "2023-06"
BINS = [-1, 0.02, 0.03, 0.05, 1]
LABELS = ["2% 미만", "2~3%", "3~5%", "5% 이상"]
BAD = -0.03                                   # "주식이 크게 빠진 달": 한 달 -3% 이하 (3% 이상 하락)
lines = []


def log(s=""):
    print(s)
    lines.append(str(s))


def load(synthetic=False):
    """SP500, Dividend, Long Interest Rate, Consumer Price Index 네 열 (월별)."""
    if synthetic:
        rng = np.random.default_rng(5)
        idx = pd.date_range("1940-01-01", END + "-01", freq="MS")
        n = len(idx)
        infl = np.clip(0.03 + np.cumsum(rng.normal(0, 0.002, n)), -0.02, 0.12)
        cpi = 20 * np.cumprod(1 + infl / 12 + rng.normal(0, 0.002, n))
        px = 25 * np.cumprod(1 + rng.normal(0.006, 0.04, n))
        y = np.clip(4 + np.cumsum(rng.normal(0, 0.15, n)), 0.5, 15)
        return pd.DataFrame({"SP500": px, "Dividend": px * 0.03, "Long Interest Rate": y,
                             "Consumer Price Index": cpi}, index=idx)
    df = pd.read_csv(URL, parse_dates=["Date"], index_col="Date")
    return df[["SP500", "Dividend", "Long Interest Rate", "Consumer Price Index"]].loc[:END]


def returns(df):
    """주식·채권 월 수익률과 직전 달까지의 12개월·36개월 물가 상승률."""
    cpi = df["Consumer Price Index"]
    d = df.loc[START:]
    p, div = d["SP500"], d["Dividend"]
    stock = (p + div / 12) / p.shift(1) - 1
    y = d["Long Interest Rate"] / 100
    c, ynew = y.shift(1), y
    bond = c * (1 - (1 + ynew) ** -10) / ynew + (1 + ynew) ** -10 - 1 + c / 12
    r = pd.DataFrame({"주식": stock, "채권": bond}).dropna()
    r["물가12"] = cpi.pct_change(12).shift(1).reindex(r.index)
    r["물가36"] = ((cpi / cpi.shift(36)) ** (1 / 3) - 1).shift(1).reindex(r.index)
    r["물가_당월"] = cpi.pct_change(12).reindex(r.index)
    return r


def regime_table(r, col):
    x = r.copy()
    x["구간"] = pd.cut(x[col], BINS, labels=LABELS, right=False)
    rows = {}
    for lab, g in x.groupby("구간", observed=True):
        bad = g[g["주식"] <= BAD]
        rows[lab] = {
            "개월": len(g),
            "상관계수": g["주식"].corr(g["채권"]),
            "주식 연변동성(%)": g["주식"].std() * np.sqrt(12) * 100,
            "채권 연변동성(%)": g["채권"].std() * np.sqrt(12) * 100,
            "주식 -3% 이하 달": len(bad),
            "그달 채권 평균(%)": bad["채권"].mean() * 100,
            "그달 채권 오른 비율(%)": (bad["채권"] > 0).mean() * 100,
        }
    t = pd.DataFrame(rows).T
    t[["개월", "주식 -3% 이하 달"]] = t[["개월", "주식 -3% 이하 달"]].astype(int)
    return t


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--synthetic", action="store_true")
    a = ap.parse_args()
    os.makedirs(OUT, exist_ok=True)
    df = load(a.synthetic)
    r = returns(df)
    pd.set_option("display.width", 200)
    pd.set_option("display.max_columns", 20)
    log(f"기간: {r.index[0]:%Y-%m} ~ {r.index[-1]:%Y-%m}, {len(r)}개월")
    log(f"전체 상관계수: {r['주식'].corr(r['채권']):.3f}")

    log("\n[1] 직전 12개월 물가 상승률 구간별")
    log(regime_table(r, "물가12").round(2).to_string())
    log("\n[2] 직전 36개월 연평균 물가 상승률 구간별")
    log(regime_table(r, "물가36").round(2).to_string())
    for lab, s, e in [("1953~1989", None, "1989"), ("1990~2023", "1990", None)]:
        log(f"\n[3] {lab}, 직전 12개월 물가 상승률 구간별")
        log(regime_table(r.loc[s:e], "물가12").round(2).to_string())

    log("\n[3-1] 구간 경계와 시차를 바꾼 상관계수 (낮은 구간부터)")
    variants = {
        "2/3/5% (본문)": ("물가12", [-1, 0.02, 0.03, 0.05, 1]),
        "2/4/6%": ("물가12", [-1, 0.02, 0.04, 0.06, 1]),
        "3% 기준 둘로": ("물가12", [-1, 0.03, 1]),
        "다섯 등분": ("물가12", None),
        "두 달 전까지 12개월, 2/3/5%": ("물가12_2", [-1, 0.02, 0.03, 0.05, 1]),
    }
    x = r.copy()
    x["물가12_2"] = x["물가12"].shift(1)
    for name, (col, bins) in variants.items():
        cut = pd.qcut(x[col], 5) if bins is None else pd.cut(x[col], bins, right=False)
        cs = [g["주식"].corr(g["채권"]) for _, g in x.groupby(cut, observed=True)]
        log(f"  {name}: " + ", ".join(f"{c:+.2f}" for c in cs))
    hi = r[r["물가12"] >= 0.05]
    lo = r[r["물가12"] < 0.02]
    log("  5% 이상인 달의 시기 분포: " + ", ".join(f"{k}년대 {v}" for k, v in (hi.index.year // 10 * 10).value_counts().sort_index().items()))
    log("  2% 미만인 달의 시기 분포: " + ", ".join(f"{k}년대 {v}" for k, v in (lo.index.year // 10 * 10).value_counts().sort_index().items()))

    log("\n[4] 시기별 상관계수와 평균 물가 상승률(12개월)")
    for s, e in [("1953", "1965"), ("1966", "1981"), ("1982", "1999"), ("2000", "2020"), ("2021", "2023")]:
        x = r.loc[s:e]
        log(f"  {x.index[0]:%Y-%m}~{x.index[-1]:%Y-%m} ({len(x)}개월): 상관 {x['주식'].corr(x['채권']):+.2f}, "
            f"물가 {x['물가_당월'].mean():.1%}")

    blocks = []
    for i in range(0, len(r) - 35, 36):
        g = r.iloc[i:i + 36]
        blocks.append((g["주식"].corr(g["채권"]), g["물가_당월"].mean()))
    b = pd.DataFrame(blocks, columns=["상관", "물가"])
    log(f"\n[5] 겹치지 않는 36개월 구간 {len(b)}개: 구간 상관계수와 구간 평균 물가의 상관 {b['상관'].corr(b['물가']):.2f}")
    rc = r["주식"].rolling(36).corr(r["채권"])
    log(f"  36개월 이동 상관계수와 그 시점 12개월 물가 상승률의 상관 {rc.corr(r['물가_당월']):.2f}")

    y22 = r.loc["2022"]
    log(f"\n[6] 2022년: 주식 {(1 + y22['주식']).prod() - 1:.1%}, 채권 {(1 + y22['채권']).prod() - 1:.1%}, "
        f"월 상관 {y22['주식'].corr(y22['채권']):.2f}")

    with open(os.path.join(OUT, "stock_bond_inflation.txt"), "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    plot(r)


def plot(r):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig = plt.figure(figsize=(11, 5.2))
    gs = fig.add_gridspec(2, 2, width_ratios=[1.7, 1], height_ratios=[1.3, 1], hspace=0.12, wspace=0.28)
    ax1 = fig.add_subplot(gs[0, 0])
    ax2 = fig.add_subplot(gs[1, 0], sharex=ax1)
    ax3 = fig.add_subplot(gs[:, 1])

    rc = r["주식"].rolling(36).corr(r["채권"])
    ax1.plot(rc.index, rc, color="#1f5fbf", linewidth=1.2)
    ax1.axhline(0, color="#333333", linewidth=0.8)
    ax1.set_ylim(-1, 1)
    ax1.set_ylabel("Correlation")
    ax1.set_title(f"Stock-bond correlation, rolling 36 months ({r.index[0]:%Y-%m} to {r.index[-1]:%Y-%m})", fontsize=10)
    ax1.tick_params(labelbottom=False)

    infl = r["물가_당월"] * 100
    ax2.fill_between(infl.index, 0, infl, color="#8c8c8c", alpha=0.6, linewidth=0)
    for level, style in [(5, "--"), (2, ":")]:
        ax2.axhline(level, color="#333333", linestyle=style, linewidth=0.8)
        ax2.text(infl.index[-1], level + 0.3, f" {level}%", va="bottom", ha="left", fontsize=8, color="#333333")
    ax2.set_ylabel("CPI inflation, YoY %")

    t = regime_table(r, "물가12")
    names = ["< 2%", "2-3%", "3-5%", ">= 5%"]
    colors = ["#1f5fbf" if v < 0 else "#b22222" for v in t["상관계수"]]
    ax3.bar(names, t["상관계수"], color=colors, width=0.6)
    for k, (v, n) in enumerate(zip(t["상관계수"], t["개월"])):
        ax3.text(k, v + (0.015 if v >= 0 else -0.015), f"{v:+.2f}\n({n} mo)", ha="center",
                 va="bottom" if v >= 0 else "top", fontsize=9)
    ax3.axhline(0, color="#333333", linewidth=0.8)
    ax3.set_ylim(-0.4, 0.5)
    ax3.set_xlabel("Inflation over the prior 12 months")
    ax3.set_ylabel("Correlation of monthly returns")
    ax3.set_title("Stock-bond correlation by inflation regime", fontsize=10)

    for ax in (ax1, ax2, ax3):
        ax.spines[["top", "right"]].set_visible(False)
        ax.grid(axis="y", color="#e6e6e6", linewidth=0.6)
        ax.set_axisbelow(True)
    fig.savefig(os.path.join(OUT, "61_stock_bond_inflation.png"), dpi=130, bbox_inches="tight")
    plt.close(fig)


if __name__ == "__main__":
    main()
