"""
코스피가 하루에 3% 넘게 떨어진 다음 날, 5거래일 뒤, 20거래일 뒤의 수익률을 모든 날과 비교한다

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

데이터: FinanceDataReader 의 코스피(KS11) 일별 종가 (1995-05 ~ 2026-10-23, 글과 같은 기간), 가격 지수(배당 제외)
"큰 하락"은 종가 기준 하루 수익률이 -3% 이하인 날이다. h일 뒤 수익률은 그날 종가에서 h거래일 뒤 종가까지다.
큰 하락은 몰려서 오므로 표본끼리 겹친다. 그래서 두 가지로 걸러 함께 보여 준다.
  - 위기마다 첫 하락: 직전 큰 하락과 20거래일 넘게 떨어진 하락만 센다.
  - 겹치지 않는 하락: 앞서 센 하락에서 20거래일 넘게 지난 하락만 센다(20일 뒤 수익률 구간이 겹치지 않음).
결과:  output/drop_rebound_results.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")
END = "2026-10-23"                                # 글의 표와 같은 기간(2026년 10월 23일까지)
DROP = -0.03
HORIZONS = [1, 5, 20]
ERAS = {"1995~2008": ("1995-01-01", "2008-12-31"), "2009~2026": ("2009-01-01", "2026-12-31"),
        "2009~2025": ("2009-01-01", "2025-12-31")}


def load(synthetic=False):
    if synthetic:
        rng = np.random.default_rng(0)
        idx = pd.bdate_range("1995-05-02", END)
        return pd.Series(1000 * np.exp(np.cumsum(rng.standard_t(4, len(idx)) * 0.012 + 0.0002)), index=idx)
    import FinanceDataReader as fdr
    return fdr.DataReader("KS11", "1995-01-01", END)["Close"].dropna()


def summary(x):
    x = x.dropna()
    return {"n": len(x), "평균": x.mean(), "중앙값": x.median(), "오른 비율": (x > 0).mean(),
            "하위10%": x.quantile(0.1), "상위10%": x.quantile(0.9)}


def episode_starts(mask, gap=20):
    """직전 큰 하락과 gap 거래일 넘게 떨어진 하락만 남긴다(위기마다 첫 하락)."""
    pos = np.flatnonzero(mask.to_numpy())
    return mask.index[[p for i, p in enumerate(pos) if i == 0 or p - pos[i - 1] > gap]]


def non_overlapping(mask, gap=20):
    """앞서 남긴 하락에서 gap 거래일 넘게 지난 하락만 남긴다(뒤 수익률 구간이 겹치지 않게)."""
    pos = np.flatnonzero(mask.to_numpy())
    keep = [pos[0]] if len(pos) else []
    for p in pos[1:]:
        if p - keep[-1] > gap:
            keep.append(p)
    return mask.index[keep]


def main(synthetic=False):
    os.makedirs(OUT, exist_ok=True)
    close = load(synthetic)
    r = close.pct_change()
    ev = r <= DROP
    fwd = {h: close.shift(-h) / close - 1 for h in HORIZONS}
    days = r.notna()
    lines = [f"기간 {close.index[0].date()} ~ {close.index[-1].date()}, {len(close)}거래일, "
             f"하루 {DROP:.0%} 이하 하락 {int(ev.sum())}일"]

    lines.append("\n[h일 뒤 수익률: 큰 하락 다음 vs 모든 날]")
    table = {}
    for h in HORIZONS:
        for lab, m in [("큰 하락 다음", ev), ("모든 날", days)]:
            table[(f"{h}일 뒤", lab)] = summary(fwd[h][m])
    t = pd.DataFrame(table).T
    fmt = t.copy()
    for col in ["평균", "중앙값", "하위10%", "상위10%"]:
        fmt[col] = (t[col] * 100).map("{:+.2f}%".format)
    fmt["오른 비율"] = (t["오른 비율"] * 100).map("{:.1f}%".format)
    fmt["n"] = t["n"].astype(int)
    lines.append(fmt.to_string())

    nxt = r.shift(-1)
    lines.append("\n[다음 날 또 크게 움직일 확률]")
    lines.append(f"다음 날 {DROP:.0%} 이하: 큰 하락 다음 {(nxt[ev] <= DROP).mean():.1%}, 모든 날 {(r[days] <= DROP).mean():.1%}")
    lines.append(f"다음 날 ±3% 이상: 큰 하락 다음 {(nxt[ev].abs() >= 0.03).mean():.1%}, 모든 날 {(r[days].abs() >= 0.03).mean():.1%}")

    lines.append("\n[시기별: 평균 / 오른 비율 (큰 하락 다음 | 모든 날)]")
    era_rows = {}
    for lab, (a, b) in ERAS.items():
        inside = (close.index >= a) & (close.index <= b)
        row = {"하락일 수": int((ev & inside).sum())}
        for h in HORIZONS:
            e, d = fwd[h][ev & inside].dropna(), fwd[h][days & inside].dropna()
            row[f"{h}일 뒤"] = f"{e.mean():+.2%} / {(e > 0).mean():.0%} | {d.mean():+.2%} / {(d > 0).mean():.0%}"
            era_rows.setdefault(h, {})[lab] = (e.mean(), d.mean())
        lines.append(f"{lab}: " + "  ".join(f"{k} {v}" for k, v in row.items()))

    lines.append("\n[하락 크기별: 평균 / 오른 비율]")
    for lab, m in [("-3%~-5%", ev & (r > -0.05)), ("-5% 이하", r <= -0.05)]:
        parts = [f"{h}일 뒤 {fwd[h][m].dropna().mean():+.2%} / {(fwd[h][m].dropna() > 0).mean():.0%}" for h in HORIZONS]
        lines.append(f"{lab} ({int(m.sum())}일): " + "  ".join(parts))

    for lab, picked in [("위기마다 첫 하락", episode_starts(ev)), ("겹치지 않는 하락", non_overlapping(ev))]:
        lines.append(f"\n[{lab}만: {len(picked)}일]  평균 / 오른 비율 / 평균의 표준오차")
        parts = []
        for h in HORIZONS:
            x = fwd[h].loc[picked].dropna()
            parts.append(f"{h}일 뒤 {x.mean():+.2%} / {(x > 0).mean():.0%} / {x.std() / np.sqrt(len(x)):.2%}")
        lines.append("  ".join(parts))

    text = "\n".join(lines)
    print(text)
    with open(os.path.join(OUT, "drop_rebound_results.txt"), "w", encoding="utf-8") as fh:
        fh.write(text + "\n")
    chart(t, era_rows)


def chart(t, era_rows):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    colors = {"큰 하락 다음": "#1f5fbf", "모든 날": "#8c8c8c"}
    names = {"큰 하락 다음": "After a ≤−3% day", "모든 날": "All days"}
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.6), gridspec_kw={"width_ratios": [3, 2]})

    for i, h in enumerate(HORIZONS):
        for j, lab in enumerate(["큰 하락 다음", "모든 날"]):
            s = t.loc[(f"{h}일 뒤", lab)]
            x = i + (-0.12 if j == 0 else 0.12)
            ax1.plot([x, x], [s["하위10%"] * 100, s["상위10%"] * 100], color=colors[lab], linewidth=2,
                     solid_capstyle="round")
            ax1.plot(x, s["중앙값"] * 100, "o", color=colors[lab], markersize=8,
                     markeredgecolor="white", markeredgewidth=1.5, label=names[lab] if i == 0 else None)
    ax1.axhline(0, color="#999999", linewidth=0.8)
    ax1.set_xticks(range(len(HORIZONS)), ["Next day", "5 days later", "20 days later"])
    ax1.set_ylabel("Return (%)")
    ax1.set_title("KOSPI: median (dot) and 10th–90th percentile (line)", fontsize=10)
    ax1.legend(frameon=False, loc="upper left")

    eras = ["1995~2008", "2009~2026"]
    xs = np.arange(len(eras))
    ev_vals = [era_rows[20][e][0] * 100 for e in eras]
    all_vals = [era_rows[20][e][1] * 100 for e in eras]
    b1 = ax2.bar(xs - 0.18, ev_vals, width=0.34, color=colors["큰 하락 다음"], label=names["큰 하락 다음"])
    b2 = ax2.bar(xs + 0.18, all_vals, width=0.34, color=colors["모든 날"], label=names["모든 날"])
    for bars in (b1, b2):
        for b in bars:
            v = b.get_height()
            ax2.annotate(f"{v:+.1f}%", (b.get_x() + b.get_width() / 2, v), ha="center",
                         va="bottom" if v >= 0 else "top", fontsize=9, color="#333333",
                         xytext=(0, 2 if v >= 0 else -2), textcoords="offset points")
    ax2.axhline(0, color="#999999", linewidth=0.8)
    ax2.set_xticks(xs, ["1995–2008", "2009–2026"])
    ax2.set_title("Mean 20-day return, by period", fontsize=10)
    ax2.legend(frameon=False, loc="upper left")
    for ax in (ax1, ax2):
        ax.spines[["top", "right"]].set_visible(False)
        ax.grid(axis="y", color="#e6e6e6", linewidth=0.6)
        ax.set_axisbelow(True)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "drop_rebound.png"), dpi=130)
    plt.close(fig)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--synthetic", action="store_true")
    main(ap.parse_args().synthetic)