"""
Growth, quality, cash-flow and valuation factors in US large caps, long-only, versus QQQ.

FinLab research, 2026-10. Reproduces every number in
https://finlab.finance/blog/us-large-cap-factors-backtest

What it does
  1. Builds point-in-time large-cap universes from US-headquartered common stocks
     (delisted names kept, one share class per company, financials excluded):
     the 200 (or 100) largest by market value, or the 500 most traded by 60-day
     dollar volume. Price > $5, > $5M a day, no +/-50% day in the past year.
  2. Scores every stock on 17 factors: growth, profitability, cash flow,
     valuation (PEG, price to sales, EV/EBITDA, earnings yield), momentum,
     low volatility, and two composites.
  3. Holds the top 20 (or 10 / 30) stocks of each factor, equal weight, rebalanced
     monthly, with finlab's default US costs (USMarket: commission 0, no tax).
  4. Compares each portfolio with QQQ buy-and-hold over the same days:
     full window, first half, second half ("passes the screen" = wins all three).

Run:  python strategy.py            (40 base portfolios, a few minutes once the data is cached)
      python strategy.py --grid     (adds the 3 x 3 grids for the three closest
                                     designs and for any design that passes)
      python strategy.py --extras   (adds the QQQ trend filter and its diagnosis,
                                     universe controls, sector-first picks,
                                     breadth timing and weekly reversal)
"""
import itertools
import json
import math
import sys
from functools import cache, lru_cache

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

START = "2017-08-01"   # first month with a year of filings and a year of prices
HOLD = 20
SECTOR_ETF = {"Technology": "XLK", "Financial Services": "XLF", "Energy": "XLE", "Healthcare": "XLV",
              "Consumer Cyclical": "XLY", "Consumer Defensive": "XLP", "Industrials": "XLI",
              "Basic Materials": "XLB", "Utilities": "XLU", "Real Estate": "XLRE", "Communication Services": "XLC"}

data.set_market("us")


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


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


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


def on_price_grid(frame):
    """Any frame -> trading day x ticker; quarterly items carry their latest filed value."""
    frame = frame.index_str_to_date()  # no-op unless the index is quarterly labels
    return frame.reindex(index=prices().index, method="ffill").reindex(columns=prices().columns)


def as_mask(frame):
    return frame.reindex(index=prices().index, columns=prices().columns).fillna(0).astype(bool)


def top(score, members, n=HOLD):
    """The n highest scores among the universe's members, as a daily True/False mask."""
    return as_mask(score[members].is_largest(n))


@cache
def sectors():
    return get("us_company_profile").drop_duplicates("stock_id").set_index("stock_id")["sector"]


# ---------------------------------------------------------------------------
# 1. Universes
# ---------------------------------------------------------------------------
@cache
def company_ids(with_financials=False):
    """US-headquartered common stocks (delisted ones included), one share class per company.

    US headquarters keeps statements and market value in the same currency
    (Canadian filers report in CAD), which every valuation ratio below needs.
    """
    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)
        & (p["country"] == "US")
    )
    if not with_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 tuple(sorted(set(rows["stock_id"]).intersection(prices().columns)))


@cache
def dollar_volume():
    close, volume = get("us_price:close"), get("us_price:volume")
    return (close * volume).rolling(60, min_periods=20).mean()


@cache
def tradable():
    """Price above $5, $5M a day of trading, and no +/-50% day in the past year."""
    daily_return = prices().pct_change(fill_method=None)
    extreme = (daily_return.abs() >= 0.5).astype(float)
    recently_extreme = extreme.rolling(252, min_periods=1).max().shift(1).fillna(0) > 0
    return as_mask((get("us_price:close") > 5) & (dollar_volume() > 5e6) & ~recently_extreme)


@cache
def universe(kind, n, with_financials=False):
    """`cap`: the n largest by market value; `dv`: the n most traded by dollar volume."""
    ids = list(company_ids(with_financials))
    size = market_value() if kind == "cap" else dollar_volume()
    size = size.reindex(columns=ids)
    return as_mask(size[tradable().reindex(columns=ids)].is_largest(n))


# ---------------------------------------------------------------------------
# 2. Statement helpers (values dated by SEC filing date)
# ---------------------------------------------------------------------------
@cache
def quarterly(key):
    quarters = get("us_balance_sheet:total_assets").index
    return get(key).reindex(index=quarters, columns=prices().columns)


def trailing_12m(key):
    return quarterly(key).rolling(4, min_periods=4).sum()


def year_over_year(ttm):
    """(this year - last year) / |last year|, only where last year was positive."""
    last_year = ttm.shift(4)
    return ((ttm - last_year) / last_year.abs())[last_year > 0]


def positive(frame):
    return frame[frame > 0]


# ---------------------------------------------------------------------------
# 3. Factors (higher score = more attractive)
# ---------------------------------------------------------------------------
def momentum_12_1():
    return prices().shift(21) / prices().shift(252) - 1


def revenue_growth():
    return year_over_year(trailing_12m("us_income_statement:revenue"))


def net_income_growth():
    return year_over_year(trailing_12m("us_income_statement:net_income"))


def gross_profitability():
    return trailing_12m("us_income_statement:gross_profit") / positive(quarterly("us_balance_sheet:total_assets"))


def roic():
    capital = quarterly("us_balance_sheet:total_debt").fillna(0) + quarterly("us_balance_sheet:total_stockholders_equity")
    return trailing_12m("us_income_statement:operating_income") / positive(capital)


def roe():
    return trailing_12m("us_income_statement:net_income") / positive(quarterly("us_balance_sheet:total_stockholders_equity"))


def operating_margin_change():
    revenue = trailing_12m("us_income_statement:revenue")
    margin = trailing_12m("us_income_statement:operating_income") / positive(revenue)
    return margin - margin.shift(4)


def rd_intensity():
    revenue = trailing_12m("us_income_statement:revenue")
    return trailing_12m("us_income_statement:research_and_development_expenses") / positive(revenue)


def fcf_yield():
    return trailing_12m("us_cash_flow:free_cash_flow") / positive(market_value())


def shareholder_yield():
    paid_out = -(trailing_12m("us_cash_flow:net_common_stock_issuance").fillna(0)
                 + trailing_12m("us_cash_flow:common_dividends_paid").fillna(0))
    return paid_out / positive(market_value())


def earnings_yield():
    return trailing_12m("us_income_statement:net_income") / positive(market_value())


def sales_yield():
    """Revenue / market value: the inverse of price to sales, so cheap scores high."""
    return trailing_12m("us_income_statement:revenue") / positive(market_value())


def low_peg():
    """PEG = P/E / (earnings growth in %); only positive earnings and growth. Lower PEG scores higher."""
    net_income = trailing_12m("us_income_statement:net_income")
    pe = market_value() / positive(net_income)
    growth = positive(year_over_year(net_income))
    return -(pe / (growth * 100))


def ebitda_to_ev():
    """EBITDA / enterprise value: the inverse of EV/EBITDA, so cheap scores high."""
    def zero(key):
        return quarterly(key).fillna(0)  # missing line item = none

    net_claims = (zero("us_balance_sheet:total_debt") + zero("us_balance_sheet:preferred_stock")
                  + zero("us_balance_sheet:minority_interest") - zero("us_balance_sheet:cash_and_short_term_investments"))
    enterprise_value = market_value() + net_claims
    return positive(trailing_12m("us_income_statement:ebitda")) / positive(enterprise_value)


def low_volatility():
    return -prices().pct_change(fill_method=None).rolling(126, min_periods=100).std()


SINGLE = {
    "net_income_growth": net_income_growth, "revenue_growth": revenue_growth,
    "operating_margin_change": operating_margin_change, "roic": roic, "roe": roe,
    "gross_profitability": gross_profitability, "rd_intensity": rd_intensity,
    "fcf_yield": fcf_yield, "shareholder_yield": shareholder_yield,
    "low_peg": low_peg, "sales_yield": sales_yield, "ebitda_to_ev": ebitda_to_ev,
    "earnings_yield": earnings_yield, "momentum_12_1": momentum_12_1, "low_volatility": low_volatility,
}
COMPOSITE = {
    "growth_momentum": ("momentum_12_1", "revenue_growth", "operating_margin_change"),
    "quality_momentum": ("momentum_12_1", "gross_profitability", "roic"),
}
NEW_IN_THIS_STUDY = ("low_peg", "sales_yield", "ebitda_to_ev", "earnings_yield", "roe")


@cache
def factor(name):
    return on_price_grid(SINGLE[name]())


@lru_cache(maxsize=1)  # consecutive designs (top 10 / 30, the grid's holds) share one universe
def composite(name, kind, size):
    """Average of the members' percentile ranks inside the universe."""
    members = universe(kind, size)
    ranks = [factor(part)[members].rank(axis=1, pct=True) for part in COMPOSITE[name]]
    return sum(ranks) / len(ranks)  # NaN unless every part is present


def portfolio(name, kind="cap", size=200, hold=HOLD):
    score = factor(name) if name in SINGLE else composite(name, kind, size)
    return top(score, universe(kind, size), hold)


# ---------------------------------------------------------------------------
# 4. Backtest and statistics
# ---------------------------------------------------------------------------
def backtest(position, resample="M", market="us", **costs):
    with data.market(market):
        report = sim(position.loc[START:], resample=resample, upload=False, **costs)
    days = prices().loc[START:].index
    return report.creturn.reindex(days).ffill().fillna(1.0)


@cache
def equal_weight(kind, size):
    """Every member of the universe, equal weight: the control for picking stocks at all."""
    return backtest(universe(kind, size))


def stats(curve):
    curve = curve.dropna()
    curve = curve / curve.iloc[0]
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    daily = curve.pct_change().dropna()
    return {
        "cagr": curve.iloc[-1] ** (1 / years) - 1,
        "sharpe": daily.mean() / daily.std() * math.sqrt(252),  # daily, no risk-free rate
        "max_drawdown": (curve / curve.cummax() - 1).min(),
        "volatility": daily.std() * math.sqrt(252),
    }


def fund(ticker, days):
    px = get("us_fund_price:adj_close")[ticker]
    px = px.reindex(days.union(px.index)).ffill().reindex(days)
    return px / px.iloc[0]


def midpoint(days):
    """The calendar midpoint of the window; both halves include this day."""
    return days[days.searchsorted(days[0] + (days[-1] - days[0]) / 2)]


def holdings(position):
    """Average holdings, share of month-ends in cash, one-way turnover per year."""
    held = position.loc[START:].resample("ME").last().astype(float)
    count = held.sum(axis=1)
    weights = held.div(count.where(count > 0), axis=0).fillna(0)
    return {
        "avg_holdings": count.mean(),
        "share_in_cash": (count == 0).mean(),
        "annual_turnover": (weights.diff().abs().sum(axis=1) / 2).iloc[1:].mean() * 12,
    }


def yearly(curve):
    year_end = curve.resample("YE").last()
    returns = year_end.pct_change()
    returns.iloc[0] = year_end.iloc[0] / curve.iloc[0] - 1
    return {str(k.year): v for k, v in returns.items()}


def sector_mix(position):
    held = position.loc[START:].resample("ME").last().astype(float)
    weights = held.div(held.sum(axis=1), axis=0)
    mix = weights.T.groupby(sectors().reindex(weights.columns).fillna("Unknown").values).sum().T.mean()
    return mix[mix > 0].sort_values(ascending=False).round(4).to_dict()


def row(curve, qqq, **extra):
    """Daily Sharpe of the strategy vs QQQ over the full window and each half; passing the screen = winning all three."""
    mid = midpoint(curve.index)
    periods = {"full": slice(None), "first_half": slice(None, mid), "second_half": slice(mid, None)}
    s = {k: stats(curve.loc[p]) for k, p in periods.items()}
    q = {k: stats(qqq.loc[p])["sharpe"] for k, p in periods.items()}
    return {**extra, "sharpe": s["full"]["sharpe"], "qqq": q["full"],
            "first_half": s["first_half"]["sharpe"], "second_half": s["second_half"]["sharpe"],
            "cagr": s["full"]["cagr"], "max_drawdown": s["full"]["max_drawdown"],
            "passes_screen": all(s[k]["sharpe"] > q[k] for k in periods)}


# ---------------------------------------------------------------------------
# 5. Designs
# ---------------------------------------------------------------------------
UNIVERSES = (("cap", 200), ("dv", 500))
BASE = (
    [(name, kind, size, HOLD) for kind, size in UNIVERSES for name in [*SINGLE, *COMPOSITE]]
    + [(name, "cap", 100, hold) for name in ("growth_momentum", "net_income_growth", "roic") for hold in (10, 30)]
)
# the three highest full-window Sharpe ratios among the base designs in the 2026-10-28 run,
# pinned so that --grid and --extras always study the article's three
CLOSEST = [("growth_momentum", "cap", 100, 10), ("net_income_growth", "cap", 200, 20), ("roe", "cap", 200, 20)]


def label(name, kind, size, hold):
    return f"{kind}{size}_{name}_top{hold}"


CLOSEST_KEYS = [label(*d) for d in CLOSEST]


def grid_for(name, kind, size, hold):
    sizes = [size // 2, size, size * 3 // 2]
    holds = [hold // 2, hold, hold * 3 // 2]
    return [(name, kind, s, h) for s, h in itertools.product(sizes, holds)]


def run_grid(design, position, curve, qqq):
    """`position` and `curve` are the design's own, already run in main(): the grid's centre cell."""
    rows = [row(curve if d == design else backtest(portfolio(*d)), qqq, universe_size=d[2], hold=d[3])
            for d in grid_for(*design)]
    return {"rows": rows, "wins": sum(r["sharpe"] > r["qqq"] for r in rows), "total": len(rows),
            "cost_0.1pct": row(backtest(position, fee_ratio=0.001, tax_ratio=0), qqq)}


# ---------------------------------------------------------------------------
# 6. Extras: the QQQ trend filter, controls and the other designs
# ---------------------------------------------------------------------------
def qqq_trend():
    """QQQ above its 200-day average and up over 6 months."""
    qqq = get("us_fund_price:adj_close")["QQQ"].reindex(prices().index).ffill()
    return (qqq > qqq.rolling(200, min_periods=100).mean()) & (qqq / qqq.shift(126) > 1)


def month_ends(days):
    return pd.DatetimeIndex(days.to_series().groupby(days.to_period("M")).last().values)


def when_on(position, on):
    """Keep the position only on days the filter is on; cash otherwise."""
    return position.mul(on.reindex(position.index).fillna(False), axis=0).astype(bool)


def filter_diagnosis(curve, qqq, on):
    """Split the unfiltered portfolio's months by what the filter said at the month-end before."""
    ends = month_ends(curve.index)
    strategy, benchmark = curve.loc[ends].pct_change().shift(-1), qqq.loc[ends].pct_change().shift(-1)
    state = on.reindex(ends).fillna(False)
    frame = pd.DataFrame({"on": state, "strategy": strategy, "qqq": benchmark}).dropna()
    out = {"months": len(frame), "months_off": int((~frame["on"]).sum())}
    for key, part in (("on", frame[frame["on"]]), ("off", frame[~frame["on"]])):
        out[key] = {"mean_monthly_strategy": part["strategy"].mean(), "mean_monthly_qqq": part["qqq"].mean(),
                    "compounded_strategy": (1 + part["strategy"]).prod() - 1,
                    "compounded_qqq": (1 + part["qqq"]).prod() - 1,
                    "share_up_months": (part["strategy"] > 0).mean()}
    spell = (frame["on"] != frame["on"].shift()).cumsum()
    spells = []
    for _, part in frame[~frame["on"]].groupby(spell[~frame["on"]]):
        spells.append({"from": str(part.index[0].date()), "months": len(part),
                       "strategy": (1 + part["strategy"]).prod() - 1, "qqq": (1 + part["qqq"]).prod() - 1})
    out["off_spells"] = spells
    return out


def daily_filter(position, on):
    """Hold the month-end picks, but go to cash the day after the filter turns off."""
    ends = month_ends(position.loc[START:].index)
    picks = position.loc[ends].reindex(position.loc[START:].index, method="ffill").fillna(False).astype(bool)
    return when_on(picks, on)


def sector_breadth(above, members):
    """Share of each sector's members trading above their 200-day average."""
    above = above[members]
    by_sector = above.T.groupby(sectors().reindex(above.columns).values).mean().T
    return by_sector.reindex(columns=list(SECTOR_ETF)).rename(columns=SECTOR_ETF)


def other_designs(qqq):
    """Sector first then stocks, breadth timing of QQQ, and weekly reversal."""
    out = {}
    wide = universe("dv", 1000, with_financials=True)
    above_200 = (prices() > prices().rolling(200, min_periods=150).mean()).astype(float)
    top_sectors = sector_breadth(above_200, wide).is_largest(3)
    etf = sectors().reindex(wide.columns).map(SECTOR_ETF).dropna()
    in_top = top_sectors[etf.values].set_axis(etf.index, axis=1)
    stock_members = wide & as_mask(in_top)
    for name in ("momentum_12_1", "roic"):
        out[f"top3_sectors_then_{name}_top20"] = row(backtest(top(factor(name), stock_members)), qqq)

    # QQQ when more than half of the 1,000 most traded stocks are above their
    # 200-day average; otherwise the strongest of IEF / SHY / GLD over 3 months.
    funds = get("us_fund_price:adj_close")
    on = above_200[wide].mean(axis=1).reindex(funds.index).ffill() > 0.5
    defensive = funds[["IEF", "SHY", "GLD"]].pct_change(63, fill_method=None).is_largest(1)
    timing = pd.DataFrame({"QQQ": on, **{t: ~on & defensive[t] for t in ("IEF", "SHY", "GLD")}})
    out["qqq_breadth_timing"] = row(backtest(FinlabDataFrame(timing), market="us_fund"), qqq)

    # Short-term reversal (Jegadeesh 1990; Lehmann 1990): the 20 worst 5-day losers, weekly.
    loser = -(prices() / prices().shift(5) - 1)
    for kind, size in UNIVERSES:
        position = top(loser, universe(kind, size))
        out[f"{kind}{size}_reversal_5d_weekly"] = row(backtest(position, resample="W"), qqq)
        out[f"{kind}{size}_reversal_5d_weekly_cost_0.1pct"] = row(
            backtest(position, resample="W", fee_ratio=0.001, tax_ratio=0), qqq)
    return out


def extras(qqq, positions, curves):
    """Returns the statistics and the daily curves of the closest designs with the filter."""
    out, filtered_curves = {}, {}
    on = qqq_trend()
    out["filter_share_on"] = float(on.loc[START:].mean())
    out["filtered"] = {}
    for key, position in positions.items():
        filtered = backtest(when_on(position, on))
        out["filtered"][key] = row(filtered, qqq)
        if key in CLOSEST_KEYS:
            filtered_curves[f"{key}_qqq_filter"] = filtered
    qqq_filtered = pd.DataFrame({"QQQ": on})
    out["qqq_filtered"] = row(backtest(FinlabDataFrame(qqq_filtered), market="us_fund"), qqq)
    out["qqq_filtered_holdings"] = holdings(qqq_filtered)
    out["filter_diagnosis"] = {key: filter_diagnosis(curves[key], qqq, on) for key in CLOSEST_KEYS}
    out["daily_filter"] = {key: row(backtest(daily_filter(positions[key], on), resample=None), qqq)
                           for key in CLOSEST_KEYS}
    out["universe_equal_weight"] = {f"{kind}{size}": row(equal_weight(kind, size), qqq)
                                    for kind, size in (("cap", 100), ("cap", 200), ("dv", 500))}
    out["other_designs"] = other_designs(qqq)
    return out, filtered_curves


def main():
    days = prices().loc[START:].index
    qqq, spy = fund("QQQ", days), fund("SPY", days)
    positions, curves, results = {}, {}, {}
    for design in BASE:
        key = label(*design)
        positions[key] = portfolio(*design)
        curves[key] = backtest(positions[key])
        results[key] = {**row(curves[key], qqq), **holdings(positions[key]),
                        "new_in_this_study": design[0] in NEW_IN_THIS_STUDY}
        r = results[key]
        print(f"{key:40s} Sharpe {r['sharpe']:.2f} vs {r['qqq']:.2f} | halves {r['first_half']:.2f} / "
              f"{r['second_half']:.2f} | CAGR {r['cagr']:.1%} MDD {r['max_drawdown']:.1%}"
              f"{' PASS' if r['passes_screen'] else ''}", flush=True)

    mid = midpoint(days)
    metrics = {
        "window": [str(days[0].date()), str(days[-1].date())], "mid": str(mid.date()),
        "qqq": stats(qqq), "qqq_halves": [stats(qqq.loc[:mid])["sharpe"], stats(qqq.loc[mid:])["sharpe"]],
        "spy": stats(spy), "qqq_yearly": yearly(qqq), "spy_yearly": yearly(spy),
        "designs": results,
        "screen_passed": [k for k, r in results.items() if r["passes_screen"]],
    }
    metrics["qqq_sharpe_to_2024"] = stats(qqq.loc[:"2024-12-31"])["sharpe"]
    for key in CLOSEST_KEYS:
        results[key].update(yearly=yearly(curves[key]), sector=sector_mix(positions[key]),
                            sharpe_to_2024=stats(curves[key].loc[:"2024-12-31"])["sharpe"])

    if "--grid" in sys.argv:
        to_test = CLOSEST + [d for d in BASE if results[label(*d)]["passes_screen"] and d not in CLOSEST]
        metrics["grid"] = {}
        for design in to_test:
            key = label(*design)
            grid = metrics["grid"][key] = run_grid(design, positions[key], curves[key], qqq)
            print(f"{key:40s} grid: {grid['wins']}/{grid['total']} beat QQQ", flush=True)
    filtered_curves = {}
    if "--extras" in sys.argv:
        metrics["extras"], filtered_curves = extras(qqq, positions, curves)

    keep = [k for k in curves if k.startswith("cap200_") or k in CLOSEST_KEYS]
    frame = pd.DataFrame({**{k: curves[k] for k in keep}, **filtered_curves})
    frame["cap200_equal_weight"] = equal_weight("cap", 200)
    frame["QQQ"], frame["SPY"] = qqq, spy
    frame.round(5).to_csv("data.csv")
    with open("metrics.json", "w") as f:
        json.dump(metrics, f, indent=1, default=float)


if __name__ == "__main__":
    main()