"""
QQQ vs SPY, VOO and VTI: buy-and-hold comparison and five QQQ / SPY switching rules
(FinLab research, 2026-11).

Holding comparison, 2016-01-04 to 2026-10-30 (FinLab fund prices start in 2016):
  dividends-included return, drawdown, and which fund won over every 1, 3 and 5 year window;
  concentration and sector mix estimated from index members and market values.

Switching rules, 2017-04-28 to 2026-10-30, each with a 3 x 3 grid around the original setting:
  relative_strength   hold whichever of QQQ / SPY has the higher L-month return
                      grid: L 9/12/15 months x check weekly/monthly/quarterly
  trend_to_spy        QQQ while it closes above its L-day average, otherwise SPY
                      grid: L 150/200/250 days x check weekly/monthly/quarterly
  trend_to_bil        QQQ while it closes above its L-day average, otherwise BIL
                      grid: as trend_to_spy
  dual_momentum       the stronger of QQQ / SPY over L months if it also beat BIL over L months,
                      otherwise BIL (relative_strength plus an absolute momentum filter)
                      grid: as relative_strength
  trend_cascade       QQQ above its average, else SPY above its own average, else BIL
                      (trend_to_spy plus trend_to_bil)
                      grid: as trend_to_spy

Signals are read at a close and traded at the next close (finlab `sim(trade_at_price="close")`).

"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 settings beat QQQ,
  2. its original setting beats QQQ in both halves of the backtest,
  3. its original setting beats QQQ at finlab's default US cost (zero) and at 0.1% per side,
  4. its original setting beats its own assets held at fixed average weights,
  5. its original setting still beats QQQ without the calendar year of its largest lead,
  6. trend_to_bil, dual_momentum and trend_cascade also beat QQQ on returns above T-bills (BIL),
  7. dual_momentum beats relative_strength and trend_to_bil, and trend_cascade beats
     trend_to_spy and trend_to_bil, on daily Sharpe at their original settings.

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

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

HOLD_START = "2016-01-04"
START = "2017-04-28"           # the first month end with 15 months of prices
END = "2026-10-30"
COST_PER_SIDE = 0.001
MONTH = 21                     # trading days
WINDOWS = {"1y": 252, "3y": 756, "5y": 1260}
WITHHOLDING = 0.30             # US tax on dividends paid to non-residents without a treaty
OUT = Path(__file__).resolve().parent
FUNDS = ["QQQ", "SPY", "VOO", "VTI"]
FREQS = ("W", "M", "Q")
MONTHS = (9, 12, 15)
AVERAGES = (150, 200, 250)

DESIGNS = {
    # name: (lookback name, lookback values, original lookback)
    "relative_strength": ("months", MONTHS, 12),
    "trend_to_spy": ("days", AVERAGES, 200),
    "trend_to_bil": ("days", AVERAGES, 200),
    "dual_momentum": ("months", MONTHS, 12),
    "trend_cascade": ("days", AVERAGES, 200),
}
HOLDS_CASH = ("trend_to_bil", "dual_momentum", "trend_cascade")
COMBINES = {"dual_momentum": ("relative_strength", "trend_to_bil"),
            "trend_cascade": ("trend_to_spy", "trend_to_bil")}

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


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


@cache
def etf_prices():
    """Fund prices with dividends, on trading days (us_fund_price also has holiday rows)."""
    prices = get("us_fund_price:adj_close")[[*FUNDS, "BIL"]]
    return prices.loc[prices["QQQ"].notna()]


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


# ---------------------------------------------------------------------------
# Weights (trading day x fund) and backtests
# ---------------------------------------------------------------------------
def momentum(months):
    prices = etf_prices()
    return prices / prices.shift(MONTH * months) - 1


def above_average(ticker, days):
    price = etf_prices()[ticker]
    return price > price.rolling(days).mean()


def weights(name, lookback):
    if name == "relative_strength":
        returns = momentum(lookback)
        qqq_leads = returns["QQQ"] > returns["SPY"]
        return pd.DataFrame({"QQQ": qqq_leads, "SPY": ~qqq_leads})
    if name in ("trend_to_spy", "trend_to_bil"):
        above = above_average("QQQ", lookback)
        fallback = "SPY" if name == "trend_to_spy" else "BIL"
        return pd.DataFrame({"QQQ": above, fallback: ~above})
    if name == "dual_momentum":
        returns = momentum(lookback)
        qqq_leads = returns["QQQ"] > returns["SPY"]
        invested = returns[["QQQ", "SPY"]].max(axis=1) > returns["BIL"]
        return pd.DataFrame({"QQQ": qqq_leads & invested, "SPY": ~qqq_leads & invested, "BIL": ~invested})
    qqq_above = above_average("QQQ", lookback)
    spy_above = above_average("SPY", lookback)
    return pd.DataFrame({"QQQ": qqq_above, "SPY": ~qqq_above & spy_above, "BIL": ~qqq_above & ~spy_above})


def position(name, lookback):
    return weights(name, lookback).astype(float).loc[START:]


def run(position, resample, fee_ratio=None):
    report = sim(position, resample=resample, trade_at_price="close", fee_ratio=fee_ratio,
                 market="us_fund", 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, lookback, resample, fee_ratio=None):
    return run(position(name, lookback), resample, fee_ratio)


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


def bench(ticker, index):
    return etf_prices()[ticker].reindex(index)


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


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


def stats(curve):
    qqq, spy = bench("QQQ", curve.index), bench("SPY", curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    returns = curve.pct_change().dropna()
    sharpe = daily_sharpe(returns)
    qqq_sharpe = daily_sharpe(qqq.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),
        "spy_sharpe": round(daily_sharpe(spy.pct_change().dropna()), 6),
        "cagr": round((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1, 6),
        "qqq_cagr": round((qqq.iloc[-1] / qqq.iloc[0]) ** (1 / years) - 1, 6),
        "spy_cagr": round((spy.iloc[-1] / spy.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((qqq / qqq.cummax() - 1).min(), 6),
        "spy_mdd": round((spy / spy.cummax() - 1).min(), 6),
        "beats_qqq": bool(sharpe > qqq_sharpe),
    }


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 without_best_year(curve):
    """Sharpe of the design and QQQ without the calendar year of the design's largest lead."""
    ours, theirs = yearly(curve), yearly(bench("QQQ", curve.index))
    best = max(ours, key=lambda year: ours[year] - theirs[year])
    returns = curve.pct_change().dropna()
    keep = returns.index.year != int(best)
    qqq = bench("QQQ", curve.index).pct_change().dropna()
    return {"year": best, "lead": round(ours[best] - theirs[best], 6),
            "sharpe": round(daily_sharpe(returns[keep]), 6), "qqq_sharpe": round(daily_sharpe(qqq[keep]), 6)}


def fixed_at_average(name):
    """The original design's funds 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),
            "above_tbills": round(excess_sharpe(curve), 6)}


def activity(name, lookback):
    """Share of month ends in each fund, and switches a year, at a monthly check."""
    w = position(name, lookback)
    held = w.loc[month_ends(w.index)].iloc[:-1]
    switches = int((held.diff().abs().sum(axis=1) > 0).sum())
    return {"share_of_months": {k: round(v, 4) for k, v in held.mean().items()},
            "switches": switches, "switches_per_year": round(switches / (len(held) / 12), 2)}


def evaluate(name):
    param, values, original = DESIGNS[name]
    grid = []
    for lookback in values:
        for check in FREQS:
            curve = backtest(name, lookback, check)
            row = {param: lookback, "check": check, **stats(curve),
                   "halves": [stats(part) for part in halves_of(curve)]}
            grid.append(row)
            print(f"{name:<18} {param}={lookback} {check}  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)
    base = next(g for g in grid if (g[param], g["check"]) == (original, "M"))
    curve = backtest(name, original, "M")
    cost = stats(backtest(name, original, "M", COST_PER_SIDE))
    wins = sum(g["beats_qqq"] for g in grid)
    fixed = fixed_at_average(name)
    dropped = without_best_year(curve)
    above_tbills = [{"sharpe": round(excess_sharpe(part), 6),
                     "qqq_sharpe": round(excess_sharpe(bench("QQQ", part.index)), 6)}
                    for part in (curve, *halves_of(curve))]
    checks = {"grid": wins >= 7, "halves": all(h["beats_qqq"] for h in base["halves"]),
              "cost": base["beats_qqq"] and cost["beats_qqq"]}
    controls = {"fixed_at_average": bool(base["sharpe"] > fixed["sharpe"]),
                "without_best_year": bool(dropped["sharpe"] > dropped["qqq_sharpe"])}
    if name in HOLDS_CASH:
        controls["above_tbills"] = bool(above_tbills[0]["sharpe"] > above_tbills[0]["qqq_sharpe"])
    return {
        "param": param, "values": list(values), "original": original, "grid": grid, "grid_wins": wins,
        "original_stats": {k: base[k] for k in base if k not in (param, "check")},
        "cost_0.1pct": cost, "fixed_at_average_weights": fixed, "without_best_year": dropped,
        "above_tbills": above_tbills, "activity": activity(name, original), "checks": checks, "controls": controls,
    }


# ---------------------------------------------------------------------------
# Holding comparison
# ---------------------------------------------------------------------------
def holding(ticker):
    curve = etf_prices()[ticker].loc[HOLD_START:]
    drawdown = curve / curve.cummax() - 1
    trough = drawdown.idxmin()
    peak = curve.loc[:trough].idxmax()
    recovered = curve.loc[trough:][curve.loc[trough:] >= curve.loc[peak]].index
    returns = curve.pct_change().dropna()
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    underwater = (drawdown < 0).astype(int)
    underwater = underwater.groupby((underwater == 0).cumsum()).cumsum()
    since_2022 = curve.loc["2022":]
    trough_2022 = since_2022.loc[:"2022"].idxmin()
    peak_2022 = curve.loc[:trough_2022].idxmax()
    after = curve.loc[trough_2022:]
    return {
        "multiple": round(curve.iloc[-1] / curve.iloc[0], 4),
        "cagr": round((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1, 6),
        "volatility": round(returns.std() * 252 ** 0.5, 6),
        "sharpe": round(daily_sharpe(returns), 6),
        "above_tbills": round(excess_sharpe(curve), 6),
        "mdd": round(drawdown.min(), 6), "mdd_peak": str(peak.date()), "mdd_trough": str(trough.date()),
        "mdd_recovered": str(recovered[0].date()) if len(recovered) else None,
        "drawdown_2022": round((curve["2022"] / curve[:"2022"].cummax().reindex(curve["2022"].index) - 1).min(), 6),
        "recovered_2022": str(after[after >= curve[peak_2022]].index[0].date()),
        "longest_underwater_days": int(underwater.max()),
        "yearly": yearly(curve),
    }


def lead_without(years_out):
    """QQQ minus SPY in return a year, without the given calendar years."""
    returns = etf_prices()[["QQQ", "SPY"]].loc[HOLD_START:].pct_change().dropna()
    kept = returns[~returns.index.year.isin(years_out)]
    growth = (1 + kept).prod() ** (252 / len(kept)) - 1
    return {"qqq": round(growth["QQQ"], 6), "spy": round(growth["SPY"], 6)}


def windows():
    prices = etf_prices()[FUNDS].loc[HOLD_START:]
    out = {}
    for label, days in WINDOWS.items():
        returns = (prices / prices.shift(days) - 1).dropna()
        leader = returns.idxmax(axis=1)
        out[label] = {
            "windows": len(returns),
            "first_start": str(prices.index[0].date()), "last_start": str(prices.index[-days - 1].date()),
            "best_of_four": {t: round((leader == t).mean(), 4) for t in FUNDS},
            "qqq_beats": {t: round((returns["QQQ"] > returns[t]).mean(), 4) for t in FUNDS[1:]},
            "voo_beats_spy": round((returns["VOO"] > returns["SPY"]).mean(), 4),
            "voo_beats_vti": round((returns["VOO"] > returns["VTI"]).mean(), 4),
            "median": {t: round(returns[t].median(), 4) for t in FUNDS},
            "worst": {t: round(returns[t].min(), 4) for t in FUNDS},
            "lost_money": {t: round((returns[t] < 0).mean(), 4) for t in FUNDS},
            "median_lead_qqq_spy": round((returns["QQQ"] - returns["SPY"]).median(), 4),
            "worst_lead_qqq_spy": round((returns["QQQ"] - returns["SPY"]).min(), 4),
        }
    return out


def rolling_lead():
    """QQQ minus SPY over every 3 year window."""
    prices = etf_prices()[FUNDS].loc[HOLD_START:]
    returns = prices / prices.shift(WINDOWS["3y"]) - 1
    return (returns["QQQ"] - returns["SPY"]).dropna()


def dividends():
    """Dividends, and return after 30% withholding on them, 2016 to the end date."""
    rows = pd.DataFrame(get("us_fund_corporate_actions"))
    rows = rows[rows["stock_id"].isin(FUNDS) & (rows["date"] <= END) & (rows["date"] >= HOLD_START)]
    close = get("us_fund_price:close")[FUNDS].reindex(etf_prices().index).loc[HOLD_START:]
    adjusted = etf_prices()[FUNDS].loc[HOLD_START:]
    paid = rows.pivot_table(index="date", columns="stock_id", values="dividend", aggfunc="sum")
    days = adjusted.index
    paid.index = days[days.searchsorted(paid.index)]
    paid = paid.groupby(level=0).sum().reindex(days).fillna(0)[FUNDS]
    years = (days[-1] - days[0]).days / 365.25
    last_year = paid.loc[days[-1] - pd.DateOffset(years=1):].iloc[1:]
    drag = (paid / close.shift(1)).fillna(0) * WITHHOLDING
    net = (adjusted.pct_change().fillna(0) - drag + 1).cumprod()
    return {t: {
        "trailing_yield": round(last_year[t].sum() / close[t].iloc[-1], 6),
        "price_cagr": round((close[t].iloc[-1] / close[t].iloc[0]) ** (1 / years) - 1, 6),
        "total_cagr": round((adjusted[t].iloc[-1] / adjusted[t].iloc[0]) ** (1 / years) - 1, 6),
        "after_withholding_cagr": round(net[t].iloc[-1] ** (1 / years) - 1, 6),
    } for t in FUNDS}


# ---------------------------------------------------------------------------
# Concentration and sectors (estimated from index members and market values)
# ---------------------------------------------------------------------------
@cache
def profile():
    return pd.DataFrame(get("us_company_profile")).astype({"stock_id": str}).drop_duplicates("stock_id").set_index("stock_id")


@cache
def market_value():
    return get("etl:us_market_value")


@cache
def members(index):
    """Trading day x ticker: True while the stock is in the index (one share class per company,
    so a company with two listed classes counts once, at its full market value)."""
    days = etf_prices().index
    m = get(f"us_index_constituents:{index}")
    names = profile()["company_name"].reindex(sorted(m.columns.intersection(market_value().columns)))
    m = m[names.index].astype(bool).reindex(days, 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
    return m


def cap_weights(index, day):
    on = members(index).loc[day]
    values = market_value().loc[day, on[on].index].dropna()
    return (values / values.sum()).sort_values(ascending=False)


def all_stock_weights(day):
    """Market-value weights of US companies' common stocks (no ADRs), one share class per company."""
    listed = get("etl:us_common_stock_filter").loc[day].fillna(False).astype(bool)
    info = profile().reindex(listed[listed].index)
    info = info[(info["country"] == "US") & ~info["is_adr"].fillna(False).astype(bool)]
    info = info[~info["company_name"].duplicated()]
    values = market_value().loc[day].reindex(info.index).dropna()
    values = values[values > 0]
    return (values / values.sum()).sort_values(ascending=False)


def concentration(w):
    return {"names": len(w), "top_1": round(w.iloc[0], 4), "top_5": round(w.iloc[:5].sum(), 4),
            "top_10": round(w.iloc[:10].sum(), 4), "effective_names": round(1 / (w ** 2).sum(), 1),
            "largest": list(w.index[:10])}


def sectors(w):
    sector = profile()["sector"].reindex(w.index).fillna("Other").replace("", "Other")
    return {k: round(v, 4) for k, v in w.groupby(sector).sum().sort_values(ascending=False).items()}


def holdings_view():
    """Weights are plain market value. S&P 500 weights use float, and the Nasdaq-100 caps its
    largest members, so both overstate the top names; read them for direction, not level."""
    days = etf_prices().index
    year_ends = [d for d in month_ends(days) if d.month == 12] + [pd.Timestamp("2026-03-31"), days[-1]]
    out = {"history": {}, "sectors": {}}
    for day in year_ends:
        out["history"][str(day.date())] = {"nasdaq100": concentration(cap_weights("nasdaq100", day)),
                                           "sp500": concentration(cap_weights("sp500", day))}
    last = days[-1]
    ndx, spx = cap_weights("nasdaq100", last), cap_weights("sp500", last)
    total = all_stock_weights(last)
    out["all_us_stocks"] = concentration(total)
    out["sp500_share_of_all_us_stocks"] = round(total.reindex(spx.index).sum(), 4)
    out["sectors"] = {"nasdaq100": sectors(ndx), "sp500": sectors(spx), "all_us_stocks": sectors(total)}
    out["sectors_2016"] = {"nasdaq100": sectors(cap_weights("nasdaq100", year_ends[0])),
                           "sp500": sectors(cap_weights("sp500", year_ends[0]))}
    shared = ndx.index.intersection(spx.index)
    out["overlap"] = {"nasdaq100_names_in_sp500": len(shared), "nasdaq100_weight_in_sp500": round(ndx[shared].sum(), 4),
                      "sp500_weight_in_nasdaq100": round(spx[shared].sum(), 4)}
    return out


def main():
    prices = etf_prices()
    out = {"window": {"hold_start": HOLD_START, "start": START, "end": END}, "cost_per_side": COST_PER_SIDE,
           "run_date": str(date.today()), "finlab_version": finlab.__version__,
           "pandas_version": pd.__version__, "designs": {}}
    for name in DESIGNS:
        out["designs"][name] = evaluate(name)
    for name, parts in COMBINES.items():
        sharpe = out["designs"][name]["original_stats"]["sharpe"]
        out["designs"][name]["controls"]["beats_singles"] = all(
            sharpe > out["designs"][part]["original_stats"]["sharpe"] for part in parts)
    for name, result in out["designs"].items():
        result["passes"] = all(result["checks"].values()) and all(result["controls"].values())
        print(name, "PASS" if result["passes"] else "FAIL", result["grid_wins"], result["checks"], result["controls"], flush=True)

    out["buy_and_hold"] = {t: holding(t) for t in FUNDS}
    out["windows"] = windows()
    out["dividends"] = dividends()
    returns = prices[FUNDS].loc[HOLD_START:].pct_change().dropna()
    out["correlation"] = {k: {j: round(v, 4) for j, v in row.items()} for k, row in returns.corr().to_dict().items()}
    out["beta_qqq_on_spy"] = round(returns["QQQ"].cov(returns["SPY"]) / returns["SPY"].var(), 4)
    out["return_without_2020_2023"] = lead_without([2020, 2023])
    out["return_without_2020_2022_2023"] = lead_without([2020, 2022, 2023])
    out["holdings"] = holdings_view()
    members_2016 = get("us_index_constituents:sp500").loc["2016-12-30"].fillna(False).astype(bool)
    out["sp500_2016"] = {"members": int(members_2016.sum()),
                         "with_market_value": out["holdings"]["history"]["2016-12-30"]["sp500"]["names"]}

    curves = {name: backtest(name, DESIGNS[name][2], "M") for name in DESIGNS}
    out["switch_window_buy_and_hold"] = {t: stats(prices[t].loc[START:]) for t in FUNDS}
    out["switch_yearly"] = {name: yearly(c) for name, c in curves.items()}
    frame = pd.DataFrame({name: c / c.iloc[0] for name, c in curves.items()})
    frame = pd.DataFrame({t: prices[t].loc[HOLD_START:] / prices[t].loc[HOLD_START] for t in FUNDS}).join(frame)
    frame["QQQ minus SPY, 3 year return"] = rolling_lead()
    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()
