"""
Four ways of holding US assets that differ from "QQQ trend on, leveraged ETF; off, defensive ETF".

FinLab research, 2026-10.

Designs (each with a 3 x 3 grid around the original setting, the original in the middle)
  long_short_momentum   market neutral: long the N stocks with the best 12-1 month return,
                        short the N worst, from the 1,000 most traded US common stocks
                        grid: lookback 9/12/15 months x N 20/30/50
  long_short_profit     market neutral: long the N stocks with the highest gross profit /
                        total assets (Novy-Marx 2013), short the N lowest, ex financials
                        grid: N 20/30/50 x universe 500/1000/1500 most traded
  large_cap_trend       the N largest US stocks by market value; each one is held (1/N of
                        the account) only while it closes above its own L-day average,
                        otherwise its slot sits in cash
                        grid: N 50/100/200 x L 150/200/250 days
  cross_asset_tsmom     time-series momentum (Moskowitz, Ooi & Pedersen 2012) on stocks
                        (QQQ), Treasuries (IEF), gold (GLD) and commodities (DBC): each asset
                        is held only while its L-month return is positive, sized by inverse
                        volatility; the rest goes to short-term Treasuries (SHY)
                        grid: lookback 9/12/15 months x volatility window 21/63/126 days
  vol_target            QQQ scaled to a target volatility (Moreira & Muir 2017) without
                        leverage: QQQ weight = min(1, target / realised volatility), the rest
                        in SHY
                        grid: target 10/15/20% x volatility window 10/21/42 days

Every design rebalances at each month end and trades at the next close
(finlab `sim(resample="M", 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 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).

Two controls show where a pass comes from:
  - Sharpe on returns above T-bills (BIL). The site's Sharpe deducts no risk-free rate,
    which flatters anything that sits in short Treasuries: SHY alone scores 1.12 against
    QQQ's 0.89 on it.
  - The same assets at their average weights, held fixed. A timing rule that cannot beat
    its own average allocation adds nothing.

large_cap_trend leaves the unused slots in cash, which earns nothing in the backtest;
a third check credits that cash with SHY's return.

Run:  python research.py      writes metrics.json and data.csv next to this file
Needs finlab 2.2.21 or later (older versions trade US funds on market holidays).
"""
import json
from functools import cache
from pathlib import Path

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

SIGNAL_START = "2017-07-31"   # first month end: a year of prices and of quarterly filings behind it
END = "2026-10-29"
COST_PER_SIDE = 0.001
UNIVERSE_SIZE = 1000
OUT = Path(__file__).resolve().parent

TSMOM_ASSETS = ["QQQ", "IEF", "GLD", "DBC"]   # stocks, Treasuries, gold, commodities
SAFE = "SHY"                                 # 1-3 year Treasuries

DESIGNS = {
    # name: (grid parameter names, grid values, original values)
    "long_short_momentum": (("months", "n"), ((9, 12, 15), (20, 30, 50)), (12, 30)),
    "long_short_profit": (("n", "universe"), ((20, 30, 50), (500, 1000, 1500)), (30, 1000)),
    "large_cap_trend": (("n", "days"), ((50, 100, 200), (150, 200, 250)), (100, 200)),
    "cross_asset_tsmom": (("months", "vol_days"), ((9, 12, 15), (21, 63, 126)), (12, 63)),
    "vol_target": (("target", "vol_days"), ((0.10, 0.15, 0.20), (10, 21, 42)), (0.15, 21)),
}

data.truncate_end = END


@cache
def get(key, market="us"):
    data.set_market(market)
    return data.get(key)


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


@cache
def trading_days():
    return prices().loc[SIGNAL_START:END].index


@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", "us_fund")[[*TSMOM_ASSETS, SAFE, "BIL"]].reindex(prices().index)


@cache
def rebalance_dates():
    """The last trading day of every month."""
    days = trading_days()
    return pd.DatetimeIndex(pd.Series(days, index=days).groupby(days.to_period("M")).max().values)


# ---------------------------------------------------------------------------
# Stock universe
# ---------------------------------------------------------------------------
@cache
def common_stocks(ex_financials):
    """Common stocks (delisted ones included), one share class per company."""
    p = get("us_company_profile")
    keep = (
        (p["security_type"] == "common_stock")
        & ~p["is_etf"].fillna(False).astype(bool)
        & ~p["is_fund"].fillna(False).astype(bool)
        & ~p["is_adr"].fillna(False).astype(bool)
    )
    if ex_financials:
        keep &= p["sector"] != "Financial Services"
    rows = p.loc[keep, ["stock_id", "company_name"]].astype(str).sort_values("stock_id")
    named = rows["company_name"].str.len() > 0
    rows = pd.concat([rows[named].drop_duplicates("company_name"), rows[~named]])
    return pd.Index(rows["stock_id"]).intersection(prices().columns)


@cache
def liquid(size, ex_financials=False):
    """The `size` most traded stocks: price > $5, 60-day average dollar volume > $5M and
    no +/-50% day in the past year (split errors, ticker reuse). Rebalance date x ticker."""
    ids = common_stocks(ex_financials)
    close, volume = get("us_price:close")[ids], get("us_price:volume")[ids]
    dollar_volume = (close * volume).rolling(60, min_periods=20).mean()
    daily_return = prices()[ids].pct_change(fill_method=None)
    recently_extreme = (daily_return.abs() >= 0.5).astype(float).rolling(252, min_periods=1).max() > 0
    ok = (close > 5) & (dollar_volume > 5e6) & ~recently_extreme
    return dollar_volume[ok].reindex(rebalance_dates()).is_largest(size)


def on_dates(frame):
    """Any frame -> rebalance date x ticker; quarterly items carry their latest filed value."""
    if len(frame.index) and isinstance(frame.index[0], str):
        frame = frame.index_str_to_date()
    frame = frame.reindex(columns=prices().columns)
    return frame.reindex(frame.index.union(rebalance_dates())).ffill().reindex(rebalance_dates())


# ---------------------------------------------------------------------------
# Weights (rebalance date x asset)
# ---------------------------------------------------------------------------
def long_short(score, universe, n):
    """+1/N on the N best scores, -1/N on the N worst; finlab scales the book to 50/50."""
    score = score.reindex(columns=universe.columns)[universe]
    top = score.rank(axis=1, ascending=False, method="first") <= n
    bottom = score.rank(axis=1, ascending=True, method="first") <= n
    return (top.astype(float) - bottom.astype(float)) / n


@cache
def momentum(months):
    p = prices()
    return on_dates(p.shift(21) / p.shift(21 * months) - 1)


@cache
def gross_profitability():
    quarters = get("us_balance_sheet:total_assets").index
    gp = get("us_income_statement:gross_profit").reindex(quarters).rolling(4).sum()
    assets = get("us_balance_sheet:total_assets").reindex(index=quarters, columns=gp.columns)
    return on_dates(gp / assets[assets > 0])


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


@cache
def weights(name, a, b):
    if name == "long_short_momentum":
        return long_short(momentum(a), liquid(UNIVERSE_SIZE), b)
    if name == "long_short_profit":
        return long_short(gross_profitability(), liquid(b, ex_financials=True), a)
    if name == "large_cap_trend":
        universe = liquid(UNIVERSE_SIZE)
        largest = market_value().reindex(columns=universe.columns)[universe].is_largest(a)
        p = prices()[largest.columns]
        above = (p > p.rolling(b).mean()).reindex(rebalance_dates())
        return (largest & above).astype(float) / a
    p = etf_prices()
    vol = lambda days: p.pct_change(fill_method=None).rolling(days).std() * 252 ** 0.5
    if name == "cross_asset_tsmom":
        assets = p[TSMOM_ASSETS]
        rising = (assets / assets.shift(21 * a) - 1 > 0).reindex(rebalance_dates())
        inverse = 1 / vol(b)[TSMOM_ASSETS].reindex(rebalance_dates())
        w = inverse.div(inverse.sum(axis=1), axis=0).where(rising, 0.0)
    elif name == "vol_target":
        w = (a / vol(b)[["QQQ"]]).clip(upper=1).reindex(rebalance_dates())
    else:
        raise KeyError(name)
    w[SAFE] = 1 - w.sum(axis=1)
    return w


def run(w, fee_ratio=None):
    market = "us_fund" if SAFE in w.columns else "us"
    position = w.reindex(trading_days(), method="ffill").fillna(0)
    data.set_market(market)
    report = sim(position, resample="M", trade_at_price="close", fee_ratio=fee_ratio, upload=False)
    # start every curve (and QQQ) on the day the first trade fills
    curve = report.creturn.loc[:END]
    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):
    return run(weights(name, a, b), fee_ratio)


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


def daily_sharpe(curve):
    return annualised_sharpe(curve.pct_change().dropna())


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 annualised_sharpe(returns - bills)


def split(curve):
    """First and second half of the backtest; both include the middle day."""
    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, *split(curve))]


def stats(curve):
    curve = curve.dropna()
    bench = qqq(curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    returns = curve.pct_change().dropna()
    sharpe, qqq_sharpe = daily_sharpe(curve), daily_sharpe(bench)
    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),
        "volatility": round(returns.std() * 252 ** 0.5, 6),
        "mdd": round((curve / curve.cummax() - 1).min(), 6),
        "beats_qqq": bool(sharpe > qqq_sharpe),
    }


def halves(curve):
    return [stats(part) for part in split(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 exposure(w):
    """Average share of the account in risky assets (long + |short| for the long-short books)."""
    risky = w.drop(columns=[SAFE], errors="ignore").iloc[:-1]
    return round(float(risky.abs().sum(axis=1).clip(upper=1).mean()), 6)


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),
                   "exposure": exposure(weights(name, a, b))}
            row["all_three_periods"] = row["beats_qqq"] and all(h["beats_qqq"] for h in row["halves"])
            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, a, b))[0] for a in va for b in vb]
    return {
        "params": [pa, pb], "values": [list(va), list(vb)], "original": [oa, ob],
        "grid": grid, "grid_wins": wins, "screen_passes": sum(g["all_three_periods"] for g in grid),
        "above_tbills": {"original_full_first_second": versus_qqq_above_tbills(backtest(name, oa, ob)),
                         "grid_wins": int(sum(x["sharpe"] > x["qqq_sharpe"] for x in above_tbills))},
        "original_stats": {k: original[k] for k in original if k not in (pa, pb)},
        "cost_0.1pct": cost,
        "checks": {"grid": wins >= 7, "halves": all(h["beats_qqq"] for h in original["halves"]),
                   "cost": cost["beats_qqq"]},
    }


@cache
def fixed_curve(name):
    """The original design's assets at their average weights, rebalanced monthly."""
    w = weights(name, *DESIGNS[name][2])
    average = w.iloc[:-1].mean()   # the last month end trades after the data ends
    return run(pd.DataFrame([average.to_numpy()] * len(w), index=w.index, columns=w.columns)), average


def sharpe_without(curve, year):
    """Daily Sharpe with one calendar year's returns left out."""
    returns = curve.dropna().pct_change().dropna()
    returns = returns[returns.index.year != year]
    return round(annualised_sharpe(returns), 6)


def fixed_at_average(name):
    curve, average = fixed_curve(name)
    return {"weights": {k: round(v, 4) for k, v in average.items() if abs(v) > 1e-9},
            **stats(curve), "halves": halves(curve), "above_tbills": versus_qqq_above_tbills(curve)}


def idle_cash_in_shy(name, a, b):
    """Credit the cash a stock design leaves idle with SHY's daily return."""
    curve = backtest(name, a, b)
    returns = curve.pct_change().dropna()
    w = weights(name, a, b)
    idle = (1 - w.abs().sum(axis=1)).shift(1).reindex(returns.index, method="ffill").fillna(0)
    shy = etf_prices()[SAFE].pct_change(fill_method=None).reindex(returns.index).fillna(0)
    total = (1 + returns + idle * shy).cumprod()
    return pd.concat([curve.iloc[:1] / curve.iloc[0], total])


def large_cap_trend_idle_in_shy():
    (pa, pb), (va, vb), (oa, ob) = DESIGNS["large_cap_trend"]
    grid = []
    for a in va:
        for b in vb:
            curve = idle_cash_in_shy("large_cap_trend", a, b)
            grid.append({pa: a, pb: b, **stats(curve), "halves": halves(curve)})
    original = next(g for g in grid if (g[pa], g[pb]) == (oa, ob))
    return {"grid": grid, "grid_wins": sum(g["beats_qqq"] for g in grid), "original": original}


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

    out["large_cap_trend_idle_in_shy"] = large_cap_trend_idle_in_shy()
    # weights set at each month end, held through the next month
    out["month_end_weights"] = {
        "vol_target_qqq": {str(d.date()): round(v, 4)
                           for d, v in weights("vol_target", *DESIGNS["vol_target"][2])["QQQ"].items()},
        "large_cap_trend_stocks": {str(d.date()): round(v, 4)
                                   for d, v in weights("large_cap_trend", *DESIGNS["large_cap_trend"][2]).sum(axis=1).items()},
    }
    out["fixed_at_average_weights"] = {name: fixed_at_average(name) for name in ("cross_asset_tsmom", "vol_target")}
    for name, c in out["fixed_at_average_weights"].items():
        print(name, "fixed at average weights", c["weights"], f"Sharpe {c['sharpe']:.3f}", flush=True)

    vt = backtest("vol_target", *DESIGNS["vol_target"][2])
    out["vol_target_without_2022"] = {"vol_target": sharpe_without(vt, 2022),
                                      "fixed_at_average_weights": sharpe_without(fixed_curve("vol_target")[0], 2022),
                                      "QQQ": sharpe_without(qqq(vt.index), 2022)}

    curves = {name: backtest(name, *DESIGNS[name][2]) for name in DESIGNS}
    curves.update({f"{name} fixed at average weights": fixed_curve(name)[0] for name in ("cross_asset_tsmom", "vol_target")})
    index = curves["vol_target"].index
    buy_and_hold = {t: etf_prices()[t].reindex(index) for t in [*TSMOM_ASSETS, SAFE]}
    curves.update({f"{t} (buy and hold)": c for t, c in buy_and_hold.items()})
    out["buy_and_hold"] = {t: {**stats(c), "above_tbills": versus_qqq_above_tbills(c)} for t, c in buy_and_hold.items()}
    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()
