"""
Quick check: does the turn-of-the-month effect (Ariel 1987; Lakonishok &
Smidt 1988; McConnell & Xu 2008) still hold after publication, and is it
strong enough to carry leverage (3x the S&P 500 for ~4 days a month, T-bills
otherwise)?

    PYTHONPATH=. uv run python -m research.lab.turn_of_month

Data: S&P 500 price index (^GSPC) from 1928 (no dividends; ~2%/yr spread
evenly, so it barely moves a 4-day vs 17-day comparison), the 3-month T-bill
from FRED DTB3 (from 1954; 0 before).

Window (fixed in the literature, not tuned): the last trading day of a month
and the first three of the next, i.e. bought at the close of day -2 and sold
at the close of day +3. The calendar is known in advance, so close fills
are executable.
  tom_1x   S&P 500 in the window, T-bills otherwise
  tom_3x   3x daily S&P 500 in the window: 3 r - 2 rf - 0.9%/yr fee (a UPRO-like
           fund), T-bills otherwise
Periods: 1928-1987 (before publication), 1988-2007 (after), 2008-2026 (after
McConnell & Xu, the most recent).
Reported only: the pre-holiday day (Ariel 1990), inferred from a weekday gap.

The rule is written to prereg.json before the run; the script refuses to
overwrite results.json. The project has not tested calendar effects before.

Result (2026-10-08): STOP. C1, C2 and C3 all fail.
  - In-window minus out-of-window excess return:
    - 1928-87: 15.8 bp a day (t 6.3); the window earned 2.1x all the excess.
    - 1988-2007: 9.2 bp (t 2.6).
    - 2008-26: 1.2 bp (t 0.27).
    - 1988-2026 pooled: t 1.86.
  - t by decade: 1950s 5.8, 1960s 3.5, 1990s 1.4, 2000s 0.7, 2010s 0.2, 2020s 0.5.
  - tom_3x 2008-26: 5.0%/yr, Sharpe 0.27, MaxDD -47%. SPX buy-and-hold: 9.3% / 0.48.
  - Pre-holiday 2008-26: +16.6 bp (t 1.9), reported only, about 9 days a year.
Interpretation: a textbook post-publication decay. Little of the premium
is left to lever.
"""

from __future__ import annotations

import json
import os
import pickle

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS
from research.lab.levered_book import stats, tbill

OUT = os.path.join(RESULTS, "turn_of_month")
PERIODS = {"1928-1987": ("1928-01-01", "1987-12-31"), "1988-2007": ("1988-01-01", "2007-12-31"),
           "2008-2026": ("2008-01-01", "2026-09-24")}
POST = ("1988-01-01", "2026-09-24")
FEE3 = 0.009

PREREG = {
    "question": "Does the turn-of-month premium survive publication and carry 3x leverage?",
    "window": "trading days -1, +1, +2, +3 around each month end (bought at the close of day -2)",
    "continue_if_all": [
        "C1 1988-2026: mean daily excess return in-window minus out-of-window, Welch t >= 2.0",
        "C2 2008-2026: the same difference > 0 with t >= 1.5",
        "C3 2008-2026: tom_3x excess Sharpe >= 0.5",
    ],
    "decision": {"CONTINUE": "preregister a TOM leverage overlay on the book (lever only in the window)",
                 "STOP": "drop calendar effects"},
}


def welch(a: pd.Series, b: pd.Series) -> tuple[float, float]:
    d = a.mean() - b.mean()
    se = np.sqrt(a.var() / len(a) + b.var() / len(b))
    return float(d), float(d / se)


def tom_mask(idx: pd.DatetimeIndex) -> pd.Series:
    """True on days -1, +1, +2, +3 of each month (return days, close-to-close)."""
    m = idx.to_period("M")
    s = pd.Series(m, index=idx)
    pos = s.groupby(s).cumcount() + 1                      # +1, +2, ... from month start
    neg = s.groupby(s).cumcount(ascending=False) + 1       # 1 on the last day
    return pd.Series((pos <= 3).to_numpy() | (neg == 1).to_numpy(), index=idx)


def main() -> None:
    os.makedirs(OUT, exist_ok=True)
    res_path = os.path.join(OUT, "results.json")
    if os.path.exists(res_path):
        raise SystemExit(f"{res_path} exists; this study does not overwrite results")
    with open(os.path.join(OUT, "prereg.json"), "w") as fh:
        json.dump(PREREG, fh, indent=2)

    with open(os.path.join("data", "research", "long_history.pkl"), "rb") as fh:
        px = pickle.load(fh)["close"]["^GSPC"].dropna()
    px = px.loc[:"2026-09-24"]
    r = px.pct_change().dropna()
    idx = r.index
    rf = tbill(idx).fillna(0.0)
    ex = r - rf / 252
    tom = tom_mask(idx)
    gap = pd.Series(idx, index=idx).shift(-1) - pd.Series(idx, index=idx)
    pre_hol = (gap > pd.Timedelta(days=1)) & (pd.Series(idx.dayofweek, index=idx) < 4)   # next day skipped mid-week
    pre_hol |= (gap > pd.Timedelta(days=3))

    tom_1x = rf / 252 + ex.where(tom, 0.0)
    tom_3x = rf / 252 + (3 * ex - FEE3 / 252).where(tom, 0.0)
    spx_3x = rf / 252 + 3 * ex - FEE3 / 252
    res = {"share_in_window": float(tom.mean()), "periods": {}}
    print(f"share of days in window {tom.mean():.2%}")
    for name, (a, b) in {**PERIODS, "1988-2026": POST}.items():
        e, t = ex.loc[a:b], tom.loc[a:b]
        d, tt = welch(e[t], e[~t])
        dh, th = welch(e[pre_hol.loc[a:b] & ~t], e[~pre_hol.loc[a:b] & ~t])
        row = {"in_bp": float(e[t].mean() * 1e4), "out_bp": float(e[~t].mean() * 1e4), "diff_bp": d * 1e4,
               "t": tt, "share_of_total_excess": float(e[t].sum() / e.sum()) if e.sum() > 0 else float("nan"),
               "pre_holiday_diff_bp": dh * 1e4, "pre_holiday_t": th}
        for k, s in (("tom_1x", tom_1x), ("tom_3x", tom_3x), ("spx_1x", r), ("spx_3x", spx_3x)):
            st = stats(s.loc[a:b], rf.loc[a:b])
            row[k] = {"cagr": float(st["cagr"]), "sharpe": float(st["sharpe"]), "mdd": float(st["mdd"]),
                      "worst_yr": float(st["worst_yr"])}
        res["periods"][name] = row
        print(f"{name}: in {row['in_bp']:5.1f} bp/d out {row['out_bp']:5.1f} diff {row['diff_bp']:5.1f} t {tt:4.2f} "
              f"| share of excess {row['share_of_total_excess']:.2f} | pre-hol diff {row['pre_holiday_diff_bp']:5.1f} t {th:4.2f}")
        for k in ("tom_1x", "tom_3x", "spx_1x", "spx_3x"):
            s = row[k]
            print(f"     {k:7s} cagr {s['cagr']*100:5.1f}% xsh {s['sharpe']:5.2f} mdd {s['mdd']*100:6.1f}% "
                  f"worst yr {s['worst_yr']*100:6.1f}%")
    by_dec = ex.groupby(idx.year // 10 * 10).apply(lambda e: welch(e[tom.loc[e.index]], e[~tom.loc[e.index]])[1])
    print("t by decade:", by_dec.round(2).to_dict())
    res["t_by_decade"] = {int(k): float(v) for k, v in by_dec.items()}
    try:
        S = pd.read_pickle(os.path.join(RESULTS, "stream_portfolio", "streams.pkl"))
        m = lambda x: (1 + x.dropna()).resample("ME").prod() - 1
        res["corr_trend"] = float(m(tom_3x.loc["2006-10":]).corr(m(S["trend"])))
        print(f"monthly corr tom_3x vs TrendRider V10: {res['corr_trend']:.2f}")
    except FileNotFoundError:
        pass

    p1, p2 = res["periods"]["1988-2026"], res["periods"]["2008-2026"]
    checks = {"C1": p1["t"] >= 2.0, "C2": p2["diff_bp"] > 0 and p2["t"] >= 1.5, "C3": p2["tom_3x"]["sharpe"] >= 0.5}
    res["verdict"] = {"checks": checks, "decision": "CONTINUE" if all(checks.values()) else "STOP"}
    print(f"checks {checks} -> {res['verdict']['decision']}")
    with open(res_path, "w") as fh:
        json.dump(res, fh, indent=2)
    pd.DataFrame({"tom_1x": tom_1x, "tom_3x": tom_3x, "in_window": tom}).to_pickle(os.path.join(OUT, "returns.pkl"))
    from research.lab.trials import append_trials
    append_trials(os.path.join(RESULTS, "trials_us.jsonl"),
                  [{"name": f"turn_of_month:{p}:{k}", **res["periods"][p][k]} for p in res["periods"]
                   for k in ("tom_1x", "tom_3x")],
                  {"study": "turn_of_month", "role": "preregistered quick check"})


if __name__ == "__main__":
    main()