"""
Moskowitz, Ooi & Pedersen(2012)의 시계열 모멘텀을 미국 주식·채권·금 세 자산으로 작게 따라 해 본다

사용법:
    pip install pandas numpy matplotlib
    python tsmom_three_assets.py
    python tsmom_three_assets.py --synthetic     # 난수 데이터로 동작 확인
    (같은 폴더에 all_weather.py, asset_allocation_basics.py, value_factor.py 가 있어야 한다)

데이터: all_weather.py 와 같은 월별 수익률. 미국 주식(실러 S&P 500, 배당 포함), 10년 만기 국채(실러 10년물 금리로 만든 근사치),
금(LBMA 월평균 가격), 무위험 수익률(파마-프렌치 공개 사본, 2020년 7월까지). 그래서 기간은 1972년 1월 ~ 2020년 7월이다.
규칙:
  - 신호: 가장 최근 한 달을 건너뛴 그 앞 12개월(t-12 ~ t-1개월) 초과수익률의 부호(논문처럼 무위험 수익률을 뺀 값). 실러·LBMA 가격과 10년물 금리가
    월평균이라 바로 이어지는 두 달 사이에 인위적인 자기상관이 생기기 때문이다. 다른 창의 결과도 [4]에 함께 출력한다.
  - 플러스면 매수, 마이너스면 공매도. 수익률은 무위험 수익률을 뺀 초과수익률로 계산한다.
  - 비중은 직전 12개월 변동성으로 나눠 자산마다 연 10% 변동성에 맞춘다(논문은 자산별 연 40%).
  - 비교 대상: 같은 비중으로 늘 매수만 하는 경우.
결과:  output/tsmom_results.txt 와 차트 PNG
"""
import argparse
import os
import sys

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
OUT = os.path.join(HERE, "output")
ASSETS = {"stock": "미국 주식", "mid": "미국 10년 국채", "gold": "금"}
TARGET_VOL = 0.10
lines = []


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


def sharpe(r):
    r = r.dropna()
    return r.mean() / r.std() * np.sqrt(12)


def build(df, lookback=12, skip=1, excess=True):
    ex = df[list(ASSETS)].sub(df["rf"], axis=0)
    base = ex if excess else df[list(ASSETS)]     # excess=False 면 무위험 수익률을 빼지 않은 수익률로 신호
    past = (1 + base).rolling(lookback).apply(np.prod, raw=True) - 1
    past = past.shift(skip)                       # 기본값: 최근 한 달을 건너뛴 그 앞 12개월(t-12 ~ t-1개월)
    w = TARGET_VOL / (ex.rolling(12).std() * np.sqrt(12))
    ts = (np.sign(past) * w).shift(1) * ex        # 이번 달 말에 정한 포지션을 다음 달에 적용
    lo = w.shift(1) * ex
    return ex, past, ts, lo


def stats(r):
    r = r.dropna()
    cum = (1 + r).cumprod()
    return {"연 평균 초과수익": r.mean() * 12, "변동성": r.std() * np.sqrt(12), "샤프": sharpe(r),
            "최대 낙폭": (cum / cum.cummax() - 1).min()}


def main(synthetic=False):
    import all_weather
    os.makedirs(OUT, exist_ok=True)
    df = all_weather.load(synthetic).dropna(subset=["rf"])
    ex, past, ts, lo = build(df)
    port, lport = ts.mean(axis=1).dropna(), lo.mean(axis=1).dropna()
    start = max(port.index[0], lport.index[0])
    port, lport = port.loc[start:], lport.loc[start:]
    log(f"기간: {port.index[0].date()} ~ {port.index[-1].date()}, {len(port)}개월")

    log("\n[1] 최근 한 달을 건너뛴 그 앞 12개월 초과수익률이 플러스였던 달과 마이너스였던 달의 다음 달 초과수익률")
    for a, name in ASSETS.items():
        d = pd.DataFrame({"p": past[a].shift(1), "n": ex[a]}).dropna()
        up, dn = d["n"][d["p"] > 0], d["n"][d["p"] <= 0]
        t = (up.mean() - dn.mean()) / np.sqrt(up.var() / len(up) + dn.var() / len(dn))
        log(f"  {name}: 플러스 뒤 {up.mean():+.2%} ({len(up)}개월) / 마이너스 뒤 {dn.mean():+.2%} ({len(dn)}개월), t값 {t:.2f}")

    log("\n[2] 샤프 비율: 시계열 모멘텀 vs 늘 매수 (같은 변동성 맞춤)")
    rows = {}
    for a, name in ASSETS.items():
        rows[name] = (sharpe(ts[a].loc[start:]), sharpe(lo[a].loc[start:]), (np.sign(past[a]).shift(1).loc[start:] > 0).mean())
        log(f"  {name}: 모멘텀 {rows[name][0]:.2f} / 늘 매수 {rows[name][1]:.2f}  (매수 쪽에 있던 달 {rows[name][2]:.0%})")
    rows["세 자산 합산"] = (sharpe(port), sharpe(lport), np.nan)
    log(f"  세 자산 합산: 모멘텀 {rows['세 자산 합산'][0]:.2f} / 늘 매수 {rows['세 자산 합산'][1]:.2f}")
    for name, r in [("모멘텀", port), ("늘 매수", lport)]:
        s = stats(r)
        log(f"  {name} 합산: " + ", ".join(f"{k} {v:.2%}" if k != "샤프" else f"{k} {v:.2f}" for k, v in s.items()))
    d = port - lport
    log(f"  합산 월 수익률 차이(모멘텀 - 늘 매수)의 t값: {d.mean() / d.std() * np.sqrt(len(d)):.2f}")
    log(f"  모멘텀 합산과 미국 주식 초과수익률의 상관: {port.corr(ex['stock'].loc[port.index]):.2f}")
    for name, legs in [("모멘텀", ts), ("늘 매수", lo)]:
        c = legs.loc[start:].corr()
        log(f"  {name} 자산끼리의 상관: 주식-국채 {c.loc['stock', 'mid']:.2f}, 주식-금 {c.loc['stock', 'gold']:.2f}, 국채-금 {c.loc['mid', 'gold']:.2f}")

    log("\n[3] 기간을 나누면 (샤프: 모멘텀 / 늘 매수, 논문 표본은 2009년까지)")
    for a0, b0 in [("1973", "1989"), ("1990", "2009"), ("2010", "2020"), ("1990", "2020")]:
        t_, l_ = ts.loc[start:].loc[a0:b0], lo.loc[start:].loc[a0:b0]
        parts = [f"{name} {sharpe(t_[a]):.2f}/{sharpe(l_[a]):.2f}" for a, name in ASSETS.items()]
        log(f"  {a0}~{b0} ({len(t_)}개월): " + ", ".join(parts)
            + f", 합산 {sharpe(t_.mean(axis=1)):.2f}/{sharpe(l_.mean(axis=1)):.2f}")

    log(f"\n[4] 신호 창에 따른 세 자산 합산 샤프: 초과수익률 신호 / 그냥 수익률 신호 (늘 매수 {sharpe(lport):.2f})")
    for lb, sk, desc in [(12, 0, "최근 달 포함 12개월 (논문과 같은 창, 월평균 가격의 인위적 자기상관이 섞일 수 있음)"),
                         (9, 1, "최근 달을 건너뛴 9개월"), (10, 1, "최근 달을 건너뛴 10개월"),
                         (11, 1, "최근 달을 건너뛴 11개월 (흔히 12-1이라 부르는 창)"),
                         (12, 1, "최근 달을 건너뛴 12개월 (이 글의 기준)"), (13, 1, "최근 달을 건너뛴 13개월")]:
        _, _, tsw, _ = build(df, lb, sk)
        _, _, tsw_r, _ = build(df, lb, sk, excess=False)
        log(f"  {desc}: {sharpe(tsw.mean(axis=1).loc[start:]):.2f} / {sharpe(tsw_r.mean(axis=1).loc[start:]):.2f}")
    _, pastr, tsr, _ = build(df, excess=False)
    log(f"  이 글의 창이지만 초과수익률 대신 수익률로 신호: {sharpe(tsr.mean(axis=1).loc[start:]):.2f} ("
        + ", ".join(f"{name} {sharpe(tsr[a].loc[start:]):.2f}, 매수 {(np.sign(pastr[a]).shift(1).loc[start:] > 0).mean():.0%}"
                    for a, name in ASSETS.items()) + ")")
    diff = (np.sign(past) != np.sign(pastr)).shift(1).loc[start:].astype(bool)
    log("    두 신호가 다른 달: " + ", ".join(f"{name} {diff[a].sum()}개월(그중 1973~1989년 {diff[a].loc[:'1989'].sum()})"
                                          for a, name in ASSETS.items()))
    log("    기간별 합산 샤프 (초과수익률 신호 / 그냥 수익률 신호): " + ", ".join(
        f"{a0}~{b0} {sharpe(ts.loc[start:].loc[a0:b0].mean(axis=1)):.2f} / {sharpe(tsr.loc[start:].loc[a0:b0].mean(axis=1)):.2f}"
        for a0, b0 in [("1973", "1989"), ("1990", "2020")]))

    log("\n[5] 미국 주식 월 초과수익률 5분위별 모멘텀 합산 평균 (논문의 '극단적인 시장에서 강하다' 확인)")
    s = ex["stock"].loc[port.index]
    q = pd.qcut(s, 5, labels=["1(최악)", "2", "3", "4", "5(최고)"])
    tab = pd.DataFrame({"주식": s, "모멘텀": port}).groupby(q, observed=True).mean() * 100
    log(tab.round(2).to_string())

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

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    labels = ["US stocks", "US 10y bond", "Gold", "3-asset combined"]
    vals = list(rows.values())
    x = np.arange(len(labels))
    fig, ax = plt.subplots(figsize=(9, 4.2))
    b1 = ax.bar(x - 0.18, [v[0] for v in vals], width=0.34, color="#1f5fbf", label="Time-series momentum (long/short)")
    b2 = ax.bar(x + 0.18, [v[1] for v in vals], width=0.34, color="#8c8c8c", label="Always long")
    for bars in (b1, b2):
        for b in bars:
            v = b.get_height()
            ax.annotate(f"{v:.2f}", (b.get_x() + b.get_width() / 2, v), ha="center", va="bottom", fontsize=9,
                        color="#333333", xytext=(0, 2), textcoords="offset points")
    ax.axhline(0, color="#999999", linewidth=0.8)
    ax.set_xticks(x, labels)
    ax.set_ylabel("Sharpe ratio")
    ax.set_title(f"Sharpe ratio, {port.index[0]:%Y-%m} to {port.index[-1]:%Y-%m} (each asset scaled to 10% vol)", fontsize=10)
    ax.legend(frameon=False, loc="upper left")
    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, "tsmom_sharpe.png"), dpi=130)
    plt.close(fig)


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