"""Quad witching and monthly options expiration on QQQ and SPY (FinLab backtest).

Quad witching day W is the third Friday of March, June, September and December, moved to
the previous trading day when the market is closed. Monthly options expiration (opex) is
the third Friday of every month, with the same holiday rule. Windows count trading days.

  Q1  Cash in the week after quad witching: cash from W+1+o to W+L+o, QQQ otherwise
  Q2  Hold only the quad witching week: QQQ from W-L+1+o to W+o, cash otherwise
  Q3  Q1 with every monthly opex instead of quad witching only
  Q4  Q2 with every monthly opex
  Grid for all four: L 4/5/6 x o -1/0/1, original setting L=5, o=0.

"Beats QQQ" means a higher daily Sharpe ratio (mean daily return / standard deviation x
sqrt(252), no risk-free rate) than buying and holding QQQ, dividends included, over the
same dates. A design passes only if
  1. at least 7 of its 9 settings beat QQQ,
  2. its original setting beats QQQ in both halves of the backtest (split at the middle date),
  3. its original setting beats QQQ at finlab's default US cost (zero) and at 0.1% per side.
Designs in the market less than half the time must also beat QQQ held at the same average
weight and beat QQQ on returns above T-bills (BIL). A design that passes all of that is then
checked without its best year against QQQ. The same rules on SPY are compared with SPY
for reference only.

Event statistics (not pass criteria): average daily QQQ and SPY returns from W-5 to W+5
around quad witching and monthly opex, and the next Taiwan trading day after each quad
witching for 0050 and the TAIEX.

Timing: the calendar is known in advance. To earn day x (close of x-1 to close of x) the
buy happens at the close of x-1; sim(resample=None, trade_at_price="close") trades a
signal at the next close, so the signal is set on x-2.

Complete the AI-assisted setup flow at https://finlab.finance/en/setup before running this file.

    python research.py      # writes metrics.json and data.csv next to this file
"""
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd
from pandas.tseries.holiday import (
    AbstractHolidayCalendar,
    GoodFriday,
    Holiday,
    USLaborDay,
    USMartinLutherKingJr,
    USMemorialDay,
    USPresidentsDay,
    USThanksgivingDay,
    nearest_workday,
    sunday_to_monday,
)

import finlab
from finlab import data
from finlab.backtest import sim

DATA_END = "2026-11-05"     # last US close used in the article; set to None for the latest data
COST_PER_SIDE = 0.001       # 0.1% per buy and per sell, for spread and slippage
OUT = Path(__file__).resolve().parent
L_GRID, O_GRID, BASE = (4, 5, 6), (-1, 0, 1), (5, 0)
OFFSETS = range(-5, 6)


class NYSEHolidays(AbstractHolidayCalendar):
    """Regular NYSE holidays, used only for trading days after the last data row."""

    rules = [
        Holiday("NewYearsDay", month=1, day=1, observance=sunday_to_monday),
        USMartinLutherKingJr,
        USPresidentsDay,
        GoodFriday,
        USMemorialDay,
        Holiday("Juneteenth", month=6, day=19, start_date="2022-01-01", observance=nearest_workday),
        Holiday("IndependenceDay", month=7, day=4, observance=nearest_workday),
        USLaborDay,
        USThanksgivingDay,
        Holiday("Christmas", month=12, day=25, observance=nearest_workday),
    ]


def cut(table):
    return table.loc[:DATA_END] if DATA_END else table


# ---------------------------------------------------------------------------
# 1. Prices and the expiration calendar
# ---------------------------------------------------------------------------
finlab.login()
data.set_market("us_fund")
prices = cut(pd.DataFrame(data.get("us_fund_price:adj_close"))[["QQQ", "SPY", "BIL"]]).dropna()
days = prices.index
future = pd.bdate_range(days[-1] + pd.Timedelta(days=1), days[-1] + pd.Timedelta(days=62))
future = future.difference(NYSEHolidays().holidays(future[0], future[-1]))
calendar = days.append(future)


def third_friday(year, month):
    first = pd.Timestamp(year, month, 1)
    return first + pd.Timedelta(days=(4 - first.weekday()) % 7 + 14)


def expirations(months):
    """Third Friday of each listed month, or the trading day before it when the market is closed."""
    out = []
    for period in pd.period_range(calendar[0], calendar[-1], freq="M"):
        if period.month in months:
            friday = third_friday(period.year, period.month)
            if calendar[0] <= friday <= calendar[-1]:
                out.append(calendar[calendar.searchsorted(friday, side="right") - 1])
    return pd.DatetimeIndex(out)


QUAD = expirations((3, 6, 9, 12))
MONTHLY = expirations(range(1, 13))


# ---------------------------------------------------------------------------
# 2. Designs and backtests
# ---------------------------------------------------------------------------
def window(events, first, last):
    """True from event+first to event+last trading days (both ends included)."""
    mask = np.zeros(len(calendar), dtype=bool)
    for p in calendar.get_indexer(events):
        mask[max(p + first, 0):max(min(p + last, len(calendar) - 1) + 1, 0)] = True
    return mask


DESIGNS = {
    "Q1": {"name": "Cash in the week after quad witching",
           "want": lambda L, o: ~window(QUAD, 1 + o, L + o)},
    "Q2": {"name": "Hold only the quad witching week",
           "want": lambda L, o: window(QUAD, -L + 1 + o, o)},
    "Q3": {"name": "Cash in the week after every monthly opex",
           "want": lambda L, o: ~window(MONTHLY, 1 + o, L + o)},
    "Q4": {"name": "Hold only monthly opex weeks",
           "want": lambda L, o: window(MONTHLY, -L + 1 + o, o)},
}


def position(want, ticker):
    """want[x] = hold over day x; the signal is set two trading days earlier."""
    return pd.DataFrame({ticker: want[2:len(days) + 2]}, index=days).astype(float)


def backtest(pos, fee_ratio=None):
    report = sim(pos, resample=None, trade_at_price="close", fee_ratio=fee_ratio, upload=False)
    curve = report.creturn.reindex(days).ffill().fillna(1.0).loc[days[1]:]   # first buy at day 1's close
    return curve / curve.iloc[0], report


def returns(curve):
    return curve.pct_change().dropna()


def sharpe(r):
    return float(r.mean() / r.std() * math.sqrt(252))


def stats(curve):
    curve = curve / curve.iloc[0]
    r = returns(curve)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return {
        "start": str(curve.index[0].date()),
        "end": str(curve.index[-1].date()),
        "cagr": round(float(curve.iloc[-1] ** (1 / years) - 1), 6),
        "sharpe": round(sharpe(r), 6),
        "mdd": round(float((curve / curve.cummax() - 1).min()), 6),
        "vol": round(float(r.std() * math.sqrt(252)), 6),
        "total_return": round(float(curve.iloc[-1] - 1), 6),
    }


def hold(ticker, index):
    return prices[ticker].reindex(index)


def versus(curve, ticker="QQQ"):
    s, b = stats(curve), stats(hold(ticker, curve.index))
    return {"strategy": s, "bench": b, "win": bool(s["sharpe"] > b["sharpe"])}


def halves(curve):
    middle = curve.index[0] + (curve.index[-1] - curve.index[0]) / 2
    return [versus(curve.loc[:middle]), versus(curve.loc[middle:])]


def above_bills(curve):
    """Daily Sharpe of returns above T-bills (BIL), strategy and QQQ."""
    bills = returns(hold("BIL", curve.index))
    s = sharpe(returns(curve) - bills)
    q = sharpe(returns(hold("QQQ", curve.index)) - bills)
    return {"strategy": round(s, 6), "qqq": round(q, 6), "win": s > q}


def fixed_weight(curve, weight):
    """QQQ at the design's average weight, the rest in cash, rebalanced daily."""
    fixed = (1 + weight * returns(hold("QQQ", curve.index))).cumprod()
    s, f = stats(curve), stats(pd.concat([pd.Series([1.0], index=curve.index[:1]), fixed]))
    return {"weight": round(weight, 6), "fixed": f, "win": bool(s["sharpe"] > f["sharpe"])}


def best_year_removed(curve):
    rs, rb = returns(curve), returns(hold("QQQ", curve.index))
    excess = (rs + 1).groupby(rs.index.year).prod() - (rb + 1).groupby(rb.index.year).prod()
    best = int(excess.idxmax())
    s, b = sharpe(rs[rs.index.year != best]), sharpe(rb[rb.index.year != best])
    return {"best_year": best, "excess_that_year": round(float(excess.max()), 6),
            "strategy_sharpe": round(s, 6), "qqq_sharpe": round(b, 6), "win": s > b}


def trades_per_year(report, curve):
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return round(len(report.get_trades()) / years, 2)


def validate(key):
    spec = DESIGNS[key]
    grid = []
    for L in L_GRID:
        for o in O_GRID:
            row = versus(backtest(position(spec["want"](L, o), "QQQ"))[0])
            grid.append({"L": L, "o": o, "sharpe": row["strategy"]["sharpe"], "cagr": row["strategy"]["cagr"],
                         "mdd": row["strategy"]["mdd"], "qqq_sharpe": row["bench"]["sharpe"], "win": row["win"]})
    pos = position(spec["want"](*BASE), "QQQ")
    weight = float(pos["QQQ"].mean())
    curve, report = backtest(pos)
    with_cost = versus(backtest(pos, COST_PER_SIDE)[0])
    spy_curve, _ = backtest(pos.rename(columns={"QQQ": "SPY"}))
    out = {
        "name": spec["name"],
        "base": {"L": BASE[0], "o": BASE[1]},
        "grid": grid,
        "grid_wins": sum(r["win"] for r in grid),
        "full": versus(curve),
        "halves": halves(curve),
        "cost_0_1pct": with_cost,
        "in_market": round(weight, 6),
        "trades_per_year": trades_per_year(report, curve),
        "best_year_removed": best_year_removed(curve),
        "spy_reference": versus(spy_curve, "SPY"),
    }
    checks = [out["grid_wins"] >= 7, all(h["win"] for h in out["halves"]), out["full"]["win"], with_cost["win"]]
    if weight < 0.5:
        out["fixed_weight"] = fixed_weight(curve, weight)
        out["above_bills"] = above_bills(curve)
        checks += [out["fixed_weight"]["win"], out["above_bills"]["win"]]
    out["passes"] = bool(all(checks))
    print(f"{key} {spec['name']}: {out['grid_wins']}/9 beat QQQ, original Sharpe "
          f"{out['full']['strategy']['sharpe']:.2f} vs QQQ {out['full']['bench']['sharpe']:.2f}, "
          f"{'passes' if out['passes'] else 'does not pass'}")
    return out, curve


# ---------------------------------------------------------------------------
# 3. Event statistics (not pass criteria)
# ---------------------------------------------------------------------------
def mean_ci(values):
    v = pd.Series(values, dtype=float).dropna()
    se = v.std() / math.sqrt(len(v))
    return {"n": len(v), "mean": round(float(v.mean()), 6), "ci_low": round(float(v.mean() - 1.96 * se), 6),
            "ci_high": round(float(v.mean() + 1.96 * se), 6), "up_ratio": round(float((v > 0).mean()), 4)}


def event_table(close, events):
    r = close.pct_change()
    pos = [p for p in close.index.get_indexer(events) if p >= 0]
    return {"events": len(pos), "first_event": str(close.index[pos[0]].date()),
            "all_days_mean": round(float(r.dropna().mean()), 6),
            "offsets": [{"offset": j, **mean_ci([r.iloc[p + j] for p in pos if 1 <= p + j < len(r)])}
                        for j in OFFSETS]}


def week_returns(close, events):
    """Five-day returns: the expiration week (W-5 close to W close) and the week after (W to W+5)."""
    pos = [p for p in close.index.get_indexer(events) if p >= 5 and p + 5 < len(close)]
    v = close.to_numpy()
    five = close.pct_change(5).dropna()
    return {"expiration_week": mean_ci([v[p] / v[p - 5] - 1 for p in pos]),
            "week_after": mean_ci([v[p + 5] / v[p] - 1 for p in pos]),
            "all_five_day_mean": round(float(five.mean()), 6)}


def taiwan_next_day():
    """First Taiwan trading day after each quad witching: 0050 and TAIEX vs all Mondays."""
    data.set_market("tw")
    tw = pd.DataFrame({
        "0050": cut(pd.DataFrame(data.get("etl:adj_close")))["0050"],
        "TAIEX": cut(pd.DataFrame(data.get("taiex_total_index:收盤指數")))["TAIEX"],
    }).dropna()
    data.set_market("us_fund")
    r = tw.pct_change().dropna().loc[str(days[0].date()):]
    after = [r.index[r.index.searchsorted(w, side="right")] for w in QUAD if w < r.index[-1] and w <= days[-1]]
    mondays = r[r.index.day_of_week == 0]
    return {
        "dates": [str(d.date()) for d in after],
        **{name: {"after_quad_witching": mean_ci(r.loc[after, name]),
                  "all_mondays": mean_ci(mondays[name]), "all_days": mean_ci(r[name])} for name in r},
    }


def main():
    results, curves = {}, {}
    for key in DESIGNS:
        results[key], curves[key] = validate(key)
    quad = QUAD[QUAD <= days[-1]]
    monthly_only = MONTHLY[(MONTHLY <= days[-1]) & ~MONTHLY.isin(QUAD)]
    events = {
        ticker: {"quad": {**event_table(prices[ticker], quad), **week_returns(prices[ticker], quad)},
                 "monthly_non_quad": {**event_table(prices[ticker], monthly_only),
                                      **week_returns(prices[ticker], monthly_only)}}
        for ticker in ("QQQ", "SPY")
    }
    moved = [str(d.date()) for d in quad if d.day_of_week != 4]
    metrics = {
        "data_end": str(days[-1].date()),
        "cost_per_side": COST_PER_SIDE,
        "calendar": {"quad_witching": len(quad), "first": str(quad[0].date()), "last": str(quad[-1].date()),
                     "moved_for_holiday": moved, "monthly_opex": int((MONTHLY <= days[-1]).sum())},
        "benchmarks": {t: stats(hold(t, days[1:])) for t in ("QQQ", "SPY")},
        "designs": results,
        "events": events,
        "taiwan_next_day": taiwan_next_day(),
    }
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    table = pd.DataFrame({f"{k}_{DESIGNS[k]['name'].lower().replace(' ', '_')}": c for k, c in curves.items()})
    for t in ("QQQ", "SPY"):
        table[f"{t}_buy_hold"] = hold(t, table.index) / hold(t, table.index).iloc[0]
    table.round(6).to_csv(OUT / "data.csv", index_label="date")


if __name__ == "__main__":
    main()
