"""
Eight more ways of holding US stocks and ETFs in nine designs (equal weight is tested on two
indices), each tested against QQQ (FinLab research, 2026-10).

Designs (each with a 3 x 3 grid around the original setting, the original in the middle)
  mega_cap              the N largest S&P 500 members by market value, equal weight
                        grid: N 5/7/10 x rebalance weekly/monthly/quarterly
  cap_power_ndx         Nasdaq 100 members, weight proportional to market value ** alpha
  cap_power_spx         S&P 500 members, same; alpha = 0 is equal weight, alpha = 1 the index
                        grid: alpha -0.25/0/0.25 x rebalance weekly/monthly/quarterly
  dividend_growth       S&P 500 members whose trailing 12-month dividend per share (split
                        adjusted) rose year over year in each of the last N years, equal weight
                        grid: N 1/2/3 years x rebalance weekly/monthly/quarterly
  high_yield            the K S&P 500 members with the highest trailing 12-month dividend yield
                        grid: K 25/50/100 x rebalance weekly/monthly/quarterly
  shareholder_yield     the K S&P 500 members with the highest (common dividends + net buybacks
                        over the last 4 quarters) / market value (Boudoukh et al. 2007)
                        grid: K 25/50/100 x rebalance weekly/monthly/quarterly
  covered_call_switch   QQQ while QQQ closes above its L-day average, otherwise QYLD
                        (the Nasdaq 100 covered call ETF)
                        grid: L 150/200/250 days x check weekly/monthly/quarterly
  factor_rotation       the k of four S&P 500 factor ETFs (SPMO momentum, SPHQ quality,
                        SPLV low volatility, RPV value) with the best L-month return, monthly
                        grid: L 9/12/15 months x k 1/2/3
  pairs                 Gatev, Goetzmann & Rouwenhorst (2006): every half year, the P same-sector
                        S&P 500 pairs whose normalised prices stayed closest over the last 252
                        days; over the next half year, when a pair's spread moves more than z
                        formation-period standard deviations, buy the laggard and short the
                        leader (1/(2P) of the account each side); close when the spread crosses
                        zero or the half year ends
                        grid: z 1.5/2/2.5 x P 10/20/30

Stock universes use each day's index members (us_index_constituents), one share class per
company. Signals are read at a close and traded at the next close
(finlab `sim(trade_at_price="close")`). Every design starts at the first month end on which
all nine of its settings have their full look-back.

"Beats QQQ" means a higher daily Sharpe ratio (no risk-free rate) than buying and holding
QQQ, dividends included, over exactly the same dates. A design passes only if
  1. at least 7 of its 9 grid settings beat QQQ,
  2. its original setting beats QQQ in both halves of the backtest, and
  3. its original setting still beats QQQ after paying 0.1% per buy and per sell
     (finlab's US default cost is zero),
and then also beats QQQ on returns above T-bills (BIL); the two timing designs
(covered_call_switch, factor_rotation) must also beat their own assets held at fixed
average weights.

Run:  python research.py      writes metrics.json and data.csv next to this file
"""
import json
from functools import cache
from pathlib import Path

import finlab
import numpy as np
import pandas as pd
from finlab import data
from finlab.backtest import sim

START = "2016-12-30"           # the first month end with a full year of prices and dividends
END = "2026-10-30"
COST_PER_SIDE = 0.001
OUT = Path(__file__).resolve().parent
FREQS = ("W", "M", "Q")

FACTOR_ETFS = ["SPMO", "SPHQ", "SPLV", "RPV"]
COVERED_CALL_ETFS = ["QYLD", "XYLD", "JEPI", "JEPQ"]
REFERENCE_ETFS = ["SPY", "RSP", "QQQE", "VIG", "DGRO", "HDV", "SCHD"]
ETFS = ["QQQ", *COVERED_CALL_ETFS, *FACTOR_ETFS, *REFERENCE_ETFS]

DESIGNS = {
    # name: (grid parameter names, grid values, original values)
    "mega_cap": (("n", "rebalance"), ((5, 7, 10), FREQS), (7, "M")),
    "cap_power_ndx": (("alpha", "rebalance"), ((-0.25, 0, 0.25), FREQS), (0, "M")),
    "cap_power_spx": (("alpha", "rebalance"), ((-0.25, 0, 0.25), FREQS), (0, "M")),
    "dividend_growth": (("years", "rebalance"), ((1, 2, 3), FREQS), (2, "M")),
    "high_yield": (("k", "rebalance"), ((25, 50, 100), FREQS), (50, "M")),
    "shareholder_yield": (("k", "rebalance"), ((25, 50, 100), FREQS), (50, "M")),
    "covered_call_switch": (("days", "rebalance"), ((150, 200, 250), FREQS), (200, "M")),
    "factor_rotation": (("months", "k"), ((9, 12, 15), (1, 2, 3)), (12, 2)),
    "pairs": (("z", "pairs"), ((1.5, 2.0, 2.5), (10, 20, 30)), (2.0, 20)),
}
TIMING = ("covered_call_switch", "factor_rotation")
# buy-and-hold ETFs shown next to each design, over the same dates as its original setting
REFERENCES = {
    "mega_cap": ["SPY"], "cap_power_ndx": ["QQQE"], "cap_power_spx": ["RSP", "SPY"],
    "dividend_growth": ["VIG", "DGRO", "SPY"], "high_yield": ["HDV", "SCHD", "SPY"],
    "shareholder_yield": ["SPY"], "covered_call_switch": ["QYLD"], "factor_rotation": FACTOR_ETFS, "pairs": [],
}
# today's "Magnificent 7", picked with hindsight: a look-ahead reference, not a design
MAGNIFICENT_7 = ["AAPL", "MSFT", "GOOG", "AMZN", "NVDA", "META", "TSLA"]

data.truncate_end = END
data.set_market("us")


@cache
def get(key):
    return data.get(key)


def prices():
    return get("us_price:adj_close")


@cache
def etf_prices():
    """ETF prices on stock-market trading days (us_fund_price also has holiday rows)."""
    return get("us_fund_price:adj_close")[[*ETFS, "BIL"]].reindex(prices().index)


def month_ends(days):
    return pd.DatetimeIndex(pd.Series(days, index=days).groupby(days.to_period("M")).max().values)


def start(name):
    if name == "dividend_growth":
        return pd.Timestamp("2019-12-31")   # three years of year-over-year dividend changes
    if name == "factor_rotation":
        lookback = etf_prices()[FACTOR_ETFS].dropna().index[21 * 15]
        ends = month_ends(prices().index)
        return ends[ends >= lookback][0]
    return pd.Timestamp(START)


# ---------------------------------------------------------------------------
# Universe
# ---------------------------------------------------------------------------
@cache
def profile():
    return pd.DataFrame(get("us_company_profile")).astype({"stock_id": str}).drop_duplicates("stock_id").set_index("stock_id")


@cache
def members(index):
    """Trading day x ticker: True while the stock is in the index (one share class per company)."""
    m = get(f"us_index_constituents:{index}")
    names = profile()["company_name"].reindex(sorted(m.columns.intersection(prices().columns)))
    m = m[names.index].astype(bool).reindex(prices().index, method="ffill").fillna(False).astype(bool)
    for tickers in names[names.astype(str).str.len() > 0].groupby(names).groups.values():
        m[tickers] &= m[tickers].cumsum(axis=1) <= 1   # renamed tickers keep their own dates
    return m


@cache
def market_value():
    return get("etl:us_market_value").reindex(index=prices().index, columns=prices().columns)


# ---------------------------------------------------------------------------
# Dividends and buybacks
# ---------------------------------------------------------------------------
def corporate_actions(column, how):
    """Trading day x ticker frame of one corporate-action column, dated on the ex-date
    (moved to the next trading day when it falls on a holiday)."""
    days = prices().index
    rows = pd.DataFrame(get("us_corporate_actions")).dropna(subset=[column])
    rows = rows[rows["date"] <= days[-1]].astype({"stock_id": str})
    rows["day"] = days[days.searchsorted(rows["date"])]
    return rows.pivot_table(index="day", columns="stock_id", values=column, aggfunc=how).reindex(days)


@cache
def split_units():
    """How many of today's shares one 2016 share has become (product of past split factors)."""
    return corporate_actions("split_factor", "prod").reindex(columns=prices().columns).fillna(1).cumprod()


@cache
def ttm_dividends():
    """Dividends per 2016 share paid over the trailing 365 days."""
    paid = corporate_actions("dividend", "sum").fillna(0) * split_units()
    return paid.rolling("365D").sum().reindex(columns=prices().columns)


def years_ago(frame, n):
    return frame.reindex(frame.index - pd.DateOffset(years=n), method="ffill").set_axis(frame.index)


@cache
def dividend_yield():
    return ttm_dividends() / split_units() / get("us_price:close").reindex(index=prices().index, columns=prices().columns)


@cache
def shareholder_yield():
    dividends = get("us_cash_flow:common_dividends_paid")
    issuance = get("us_cash_flow:net_common_stock_issuance")
    paid_out = -(dividends.fillna(0) + issuance.fillna(0))
    paid_out = paid_out[dividends.notna() | issuance.notna()].rolling(4).sum()
    dated = paid_out.index_str_to_date()
    dated = dated.reindex(dated.index.union(prices().index)).ffill().reindex(prices().index)
    return dated.reindex(columns=prices().columns) / market_value()


# ---------------------------------------------------------------------------
# Pairs
# ---------------------------------------------------------------------------
def formation_dates():
    days = prices().loc[START:].index
    ends = month_ends(days)
    return ends[ends.month.isin([6, 12])]


@cache
def closest_pairs(f):
    """The 30 same-sector pairs with the smallest sum of squared differences in normalised
    price over the 252 days ending on f: (a, b, base price a, base price b, spread std)."""
    p = prices()
    i = p.index.get_loc(f)
    on = members("sp500").loc[f]
    window = p.iloc[i - 251:i + 1][on[on].index].dropna(axis=1)
    sector = profile()["sector"].reindex(window.columns).dropna()
    norm = window / window.iloc[0]
    found = []
    for names in sector.groupby(sector).groups.values():
        x = norm[names].to_numpy()
        sq = (x ** 2).sum(axis=0)
        distance = sq[:, None] + sq[None, :] - 2 * x.T @ x
        for a, b in zip(*np.triu_indices(len(names), 1)):
            found.append((distance[a, b], names[a], names[b]))
    found.sort()
    return [(a, b, window[a].iloc[0], window[b].iloc[0], float((norm[a] - norm[b]).std()))
            for _, a, b in found[:30]]


def pairs_weights(z, n):
    p = prices()
    weights = {}
    dates = formation_dates()
    for f, following in zip(dates, [*dates[1:], pd.Timestamp.max]):
        i = p.index.get_loc(f)
        period = p.index[i + 1:i + 127]
        period = period[period < following]
        for a, b, base_a, base_b, sigma in closest_pairs(f)[:n]:
            spread = (p.loc[period, a] / base_a - p.loc[period, b] / base_b).to_numpy()
            side, sides = 0, np.zeros(len(period))
            for t, s in enumerate(spread):
                if np.isnan(s):
                    break                       # one of the two stopped trading: stay out
                if side == 0 and abs(s) > z * sigma:
                    side = np.sign(s)           # +1: a is rich, short a and buy b
                elif side != 0 and s * side <= 0:
                    side = 0
                sides[t] = side
            leg = pd.Series(sides / (2 * n), index=period).reindex(p.index, fill_value=0)
            weights[a] = weights.get(a, 0) - leg
            weights[b] = weights.get(b, 0) + leg
    return pd.DataFrame(weights)


# ---------------------------------------------------------------------------
# Weights (trading day x asset) and backtests
# ---------------------------------------------------------------------------
def within(frame, mask):
    """frame where mask is True, blank elsewhere (aligned to the mask)."""
    return frame.reindex(index=mask.index, columns=mask.columns).where(mask)


def equal_weight(chosen):
    chosen = chosen.astype(float)
    return chosen.div(chosen.sum(axis=1), axis=0)


def top(score, k):
    return score.rank(axis=1, ascending=False, method="first") <= k


def weights(name, a, b):
    if name == "mega_cap":
        return equal_weight(top(within(market_value(), members("sp500")), a))
    if name.startswith("cap_power"):
        mv = within(market_value(), members("nasdaq100" if name.endswith("ndx") else "sp500"))
        w = (mv ** a).where(mv > 0)
        return w.div(w.sum(axis=1), axis=0)
    if name == "dividend_growth":
        ttm = ttm_dividends()[members("sp500").columns]
        rose = years_ago(ttm, a) > 0
        for year in range(a):
            rose &= years_ago(ttm, year) > years_ago(ttm, year + 1)
        return equal_weight(rose & members("sp500"))
    if name == "high_yield":
        return equal_weight(top(within(dividend_yield(), members("sp500")), a))
    if name == "shareholder_yield":
        return equal_weight(top(within(shareholder_yield(), members("sp500")), a))
    if name == "covered_call_switch":
        qqq = etf_prices()["QQQ"]
        above = (qqq > qqq.rolling(a).mean()).astype(float)
        return pd.DataFrame({"QQQ": above, "QYLD": 1 - above})
    if name == "factor_rotation":
        p = etf_prices()[FACTOR_ETFS]
        return equal_weight(top(p / p.shift(21 * a) - 1, b))
    return pairs_weights(a, b)


def position(name, a, b):
    return weights(name, a, b).loc[start(name):].fillna(0)


def run(position, resample, fee_ratio=None):
    position = position.loc[:, (position != 0).any()]
    market = "us_fund" if position.columns.isin(ETFS).all() else "us"
    report = sim(position, resample=resample, trade_at_price="close", fee_ratio=fee_ratio, market=market, upload=False)
    # start every curve (and QQQ) on the day the first trade fills
    curve = report.creturn
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):]


@cache
def backtest(name, a, b, fee_ratio=None):
    # pairs positions change on the days spreads open and close; factor_rotation's b is k
    resample = {"pairs": None, "factor_rotation": "M"}.get(name, b)
    return run(position(name, a, b), resample, fee_ratio)


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------
def daily_sharpe(returns):
    return returns.mean() / returns.std() * 252 ** 0.5


def qqq(index):
    return etf_prices()["QQQ"].reindex(index)


def excess_sharpe(curve):
    """Daily Sharpe of the returns above T-bills (BIL)."""
    returns = curve.pct_change().dropna()
    bills = etf_prices()["BIL"].reindex(curve.index).pct_change(fill_method=None).reindex(returns.index)
    return daily_sharpe(returns - bills)


def halves_of(curve):
    middle = curve.index[len(curve) // 2]
    return curve.loc[:middle], curve.loc[middle:]


def versus_qqq_above_tbills(curve):
    return [{"sharpe": round(excess_sharpe(part), 6), "qqq_sharpe": round(excess_sharpe(qqq(part.index)), 6)}
            for part in (curve, *halves_of(curve))]


def stats(curve):
    bench = qqq(curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    returns = curve.pct_change().dropna()
    sharpe, qqq_sharpe = daily_sharpe(returns), daily_sharpe(bench.pct_change().dropna())
    return {
        "start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
        "sharpe": round(sharpe, 6), "qqq_sharpe": round(qqq_sharpe, 6),
        "cagr": round((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1, 6),
        "qqq_cagr": round((bench.iloc[-1] / bench.iloc[0]) ** (1 / years) - 1, 6),
        "volatility": round(returns.std() * 252 ** 0.5, 6),
        "mdd": round((curve / curve.cummax() - 1).min(), 6),
        "qqq_mdd": round((bench / bench.cummax() - 1).min(), 6),
        "beats_qqq": bool(sharpe > qqq_sharpe),
    }


def halves(curve):
    return [stats(part) for part in halves_of(curve)]


def yearly(curve):
    ends = curve.resample("YE").last()
    returns = ends / ends.shift(1, fill_value=curve.iloc[0]) - 1
    return {str(d.year): round(v, 6) for d, v in returns.items()}


def holdings(name, a, b):
    """Average number of assets held (long or short) on the days the design trades."""
    return round(float((position(name, a, b) != 0).sum(axis=1).mean()), 2)


def evaluate(name):
    (pa, pb), (va, vb), (oa, ob) = DESIGNS[name]
    grid = []
    for a in va:
        for b in vb:
            curve = backtest(name, a, b)
            row = {pa: a, pb: b, **stats(curve), "halves": halves(curve), "holdings": holdings(name, a, b)}
            grid.append(row)
            print(f"{name:<20} {pa}={a} {pb}={b}  Sharpe {row['sharpe']:.3f} vs QQQ {row['qqq_sharpe']:.3f}"
                  f"  halves {[round(h['sharpe'], 2) for h in row['halves']]}"
                  f" vs {[round(h['qqq_sharpe'], 2) for h in row['halves']]}", flush=True)
    original = next(g for g in grid if (g[pa], g[pb]) == (oa, ob))
    cost = stats(backtest(name, oa, ob, COST_PER_SIDE))
    wins = sum(g["beats_qqq"] for g in grid)
    above_tbills = versus_qqq_above_tbills(backtest(name, oa, ob))
    checks = {"grid": wins >= 7, "halves": all(h["beats_qqq"] for h in original["halves"]),
              "cost": cost["beats_qqq"]}
    controls = {"above_tbills": bool(above_tbills[0]["sharpe"] > above_tbills[0]["qqq_sharpe"])}
    result = {
        "params": [pa, pb], "values": [list(va), list(vb)], "original": [oa, ob],
        "grid": grid, "grid_wins": wins,
        "original_stats": {k: original[k] for k in original if k not in (pa, pb)},
        "cost_0.1pct": cost, "above_tbills": above_tbills, "checks": checks,
    }
    if name in TIMING:
        result["fixed_at_average_weights"] = fixed_at_average(name)
        controls["fixed_at_average"] = bool(original["sharpe"] > result["fixed_at_average_weights"]["sharpe"])
    result["controls"] = controls
    result["passes"] = all(checks.values()) and all(controls.values())
    return result


def fixed_at_average(name):
    """The original design's assets at their average weights, rebalanced monthly."""
    w = position(name, *DESIGNS[name][2])
    average = w.loc[month_ends(w.index)].iloc[:-1].mean()
    curve = run(pd.DataFrame(dict(average), index=w.index), "M")
    return {"weights": {k: round(v, 4) for k, v in average.items()},
            **stats(curve), "halves": halves(curve), "above_tbills": versus_qqq_above_tbills(curve)}


def buy_and_hold(ticker, since):
    curve = etf_prices()[ticker].loc[since:].dropna()
    return {**stats(curve), "halves": halves(curve), "above_tbills": versus_qqq_above_tbills(curve)}


def mega_cap_months_held():
    """How many month ends each stock was among the 7 largest."""
    w = position("mega_cap", *DESIGNS["mega_cap"][2])
    held = (w.loc[month_ends(w.index)] > 0).sum()
    return {k: int(v) for k, v in held[held > 0].sort_values(ascending=False).items()}


def universe(index):
    m = get(f"us_index_constituents:{index}")
    ever = m.columns[m.any()]
    used = members(index)
    return {"ever_members": len(ever),
            "ever_members_without_prices": int((~ever.isin(prices().columns)).sum()),
            "ever_members_used": int(used.any().sum()),
            "members_on_start": int(used.loc[START].sum()),
            "with_market_value_on_start": int(within(market_value(), used).loc[START].notna().sum())}


def main():
    out = {"window": {"start": START, "end": END}, "cost_per_side": COST_PER_SIDE,
           "finlab_version": finlab.__version__, "designs": {}}
    for name in DESIGNS:
        result = evaluate(name)
        out["designs"][name] = result
        print(name, "PASS" if result["passes"] else "FAIL", result["checks"], result["controls"], flush=True)

    out["starts"] = {name: str(start(name).date()) for name in DESIGNS}
    out["buy_and_hold"] = {t: buy_and_hold(t, START) for t in ETFS}
    out["covered_call_since_listing"] = {t: buy_and_hold(t, month_ends(etf_prices()[t].dropna().index)[0])
                                         for t in COVERED_CALL_ETFS}
    curves = {name: backtest(name, *DESIGNS[name][2]) for name in DESIGNS}
    out["references"] = {name: {t: stats(etf_prices()[t].reindex(curves[name].index)) for t in REFERENCES[name]}
                         for name in DESIGNS}
    days = prices().loc[START:].index
    mag7 = run(pd.DataFrame(1 / 7, index=days, columns=MAGNIFICENT_7), "M")
    out["hindsight_magnificent_7"] = {"tickers": MAGNIFICENT_7, **stats(mag7), "halves": halves(mag7)}
    out["mega_cap_months_held"] = mega_cap_months_held()
    out["universe"] = {index: universe(index) for index in ("sp500", "nasdaq100")}

    curves["Magnificent 7 picked with hindsight"] = mag7
    curves.update({f"{t} (buy and hold)": etf_prices()[t].reindex(days) for t in ETFS})
    out["yearly"] = {name: yearly(c.dropna()) for name, c in curves.items()}
    frame = pd.DataFrame(curves)
    frame = frame / frame.bfill().iloc[0]
    frame.round(6).to_csv(OUT / "data.csv", index_label="date")
    (OUT / "metrics.json").write_text(json.dumps(out, indent=1, ensure_ascii=False))


if __name__ == "__main__":
    main()
