"""
Six famous stock screens on US equities, long-only, versus QQQ (total return).

FinLab research, 2026-10. Reproduces every number in
https://finlab.finance/blog/us-famous-stock-screens-backtest

Screens
  1. Rule of 40      revenue growth + free-cash-flow margin >= 40%, software companies
  2. Magic Formula   Greenblatt (2005): earnings yield (EBIT/EV) + return on capital ranks
  3. CANSLIM         O'Neil: quarterly and annual EPS growth, near a new high, relative
                     strength leader, market in an uptrend (institutional sponsorship and
                     share supply are left out: no point-in-time data)
  4. Graham number   price below sqrt(22.5 x EPS x book value per share)
  5. Dividend growth trailing dividends per share raised every year for 3 years
  6. Equal weight    every S&P 500 member at the same weight, rebalanced quarterly

Every portfolio is rebuilt from point-in-time data: statement items are dated by
their SEC filing date, the S&P 500 list is the membership on each day, and the
company list keeps delisted names. finlab sim() runs with its default US costs
(USMarket: commission 0, no transfer tax); a 0.1%-per-side cost run is added.

Run:  python research.py            (six base portfolios, a couple of minutes)
      python research.py --grid     (adds the neighbouring-parameter grids)
      python research.py --extras   (adds the universe controls, variants,
                                     dividend start-date and streak tests)
"""
import itertools
import json
import math
import sys
from functools import cache

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

START = "2017-08-01"          # first month with a year of filings and prices
DIVIDEND_START = "2021-01-04"  # first day a 4-year dividend record exists (data starts 2016)
ETFS = ["QQQ", "SPY", "RSP", "VIG", "SCHD", "DGRO", "DGRW", "SDY"]

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."""
    if len(frame.index) and isinstance(frame.index[0], str):
        frame = frame.index_str_to_date()
    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)


# ---------------------------------------------------------------------------
# 1. Universes
# ---------------------------------------------------------------------------
def company_ids(us_only=False, ex_sectors=(), industries=None):
    """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 us_only:
        keep &= p["country"] == "US"
    if ex_sectors:
        keep &= ~p["sector"].isin(list(ex_sectors))
    if industries:
        keep &= p["industry"].isin(list(industries))
    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 sorted(set(rows["stock_id"]).intersection(prices().columns))


@cache
def tradable():
    """Price above $5, $5M a day of trading, and no +/-50% day in the past year."""
    close, volume = get("us_price:close"), get("us_price:volume")
    dollar_volume = (close * volume).rolling(60, min_periods=20).mean()
    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((close > 5) & (dollar_volume > 5e6) & ~recently_extreme)


def largest(ids, n):
    """The n largest tradable companies by market value on each day."""
    cap = market_value().reindex(columns=ids)
    cap = cap[tradable().reindex(columns=ids)]
    return as_mask(cap.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()


# ---------------------------------------------------------------------------
# 3. The six screens (each returns a daily True/False holding table)
# ---------------------------------------------------------------------------
SOFTWARE = ["Software - Application", "Software - Infrastructure"]


@cache
def rule_of_40_score():
    revenue = trailing_12m("us_income_statement:revenue")
    revenue_year_ago = revenue.shift(4)
    growth = revenue / revenue_year_ago[revenue_year_ago > 0] - 1
    fcf_margin = trailing_12m("us_cash_flow:free_cash_flow") / revenue[revenue > 0]
    return on_price_grid(growth + fcf_margin)


def rule_of_40_universe(min_cap=2e9):
    """Tradable software companies worth at least `min_cap`."""
    ids = company_ids(industries=SOFTWARE)
    cap = market_value().reindex(columns=ids)
    return as_mask((cap >= min_cap) & tradable().reindex(columns=ids))


def rule_of_40(threshold=0.40, min_cap=2e9):
    return as_mask(rule_of_40_score() >= threshold) & rule_of_40_universe(min_cap)


@cache
def magic_formula_scores():
    zero = lambda key: quarterly(key).fillna(0)  # noqa: E731  missing line item = none
    ebit = trailing_12m("us_income_statement:ebit")
    cash = zero("us_balance_sheet:cash_and_short_term_investments")
    net_claims = (
        zero("us_balance_sheet:total_debt")
        + zero("us_balance_sheet:preferred_stock")
        + zero("us_balance_sheet:minority_interest")
        - cash
    )
    enterprise_value = market_value() + net_claims
    earnings_yield = ebit / enterprise_value[enterprise_value > 0]
    working_capital = (quarterly("us_balance_sheet:total_current_assets") - cash) - (
        quarterly("us_balance_sheet:total_current_liabilities") - zero("us_balance_sheet:short_term_debt")
    )
    capital = working_capital.clip(lower=0) + quarterly("us_balance_sheet:property_plant_equipment_net")
    return_on_capital = ebit / capital[capital > 0]
    return on_price_grid(earnings_yield), on_price_grid(return_on_capital)


def magic_formula_universe(top=1000):
    """The `top` largest US companies, ex financials, real estate, utilities."""
    return largest(company_ids(us_only=True, ex_sectors=("Financial Services", "Real Estate", "Utilities")), top)


def magic_formula(n=30, top=1000):
    """Greenblatt's ranking: earnings yield rank + return on capital rank."""
    universe = magic_formula_universe(top)
    earnings_yield, return_on_capital = magic_formula_scores()
    combined = earnings_yield[universe].rank(axis=1, pct=True) + return_on_capital[universe].rank(axis=1, pct=True)
    return as_mask(combined.is_largest(n)) & universe


@cache
def canslim_inputs():
    eps = quarterly("us_income_statement:eps_diluted")
    eps_year_ago = eps.shift(4)
    quarter_growth = (eps / eps_year_ago[eps_year_ago > 0] - 1)[eps > 0]          # C
    eps_12m = eps.rolling(4, min_periods=4).sum()
    eps_12m_year_ago = eps_12m.shift(4)
    annual_growth = eps_12m / eps_12m_year_ago[eps_12m_year_ago > 0] - 1           # A
    price = prices()
    to_high = price / price.rolling(252, min_periods=200).max()                      # N
    strength = sum(w * (price / price.shift(d) - 1)                                   # L
                   for w, d in ((0.4, 63), (0.2, 126), (0.2, 189), (0.2, 252)))
    spy = get("us_fund_price:adj_close")["SPY"].reindex(price.index).ffill()
    market_up = spy > spy.rolling(200).mean()                                         # M
    close, volume = get("us_price:close"), get("us_price:volume")
    liquid = (close >= 15) & (volume.rolling(50, min_periods=40).mean() >= 400_000)
    ids = company_ids()
    base = as_mask(liquid.reindex(columns=ids) & tradable().reindex(columns=ids))
    rs_rank = strength[base].rank(axis=1, pct=True)
    return on_price_grid(quarter_growth), on_price_grid(annual_growth), to_high, rs_rank, market_up, base


def canslim(growth=0.25, rs_min=0.80, near_high=0.85, market_filter=True):
    quarter_growth, annual_growth, to_high, rs_rank, market_up, base = canslim_inputs()
    picks = as_mask(
        (quarter_growth >= growth) & (annual_growth >= growth) & (to_high >= near_high) & (rs_rank >= rs_min)
    ) & base
    if not market_filter:
        return picks
    up = market_up.reindex(picks.index).fillna(False).astype(bool)
    return picks.mul(up, axis=0).astype(bool)


@cache
def graham_product():
    """P/E x P/B from market value, so share counts and splits cancel out."""
    net_income = trailing_12m("us_income_statement:net_income")
    equity = quarterly("us_balance_sheet:total_stockholders_equity")
    pe = market_value() / net_income[net_income > 0]
    pb = market_value() / equity[equity > 0]
    return on_price_grid(pe * pb)


def graham_universe(top=1000):
    """The `top` largest US companies, ex financials and real estate."""
    return largest(company_ids(us_only=True, ex_sectors=("Financial Services", "Real Estate")), top)


def graham_number(multiple=22.5, top=1000):
    """Price < sqrt(multiple x EPS x BVPS)  <=>  P/E x P/B < multiple."""
    return as_mask(graham_product() <= multiple) & graham_universe(top)


@cache
def dividends_12m():
    """Split-adjusted trailing-365-day dividends per share; special dividends dropped."""
    actions = pd.DataFrame(get("us_corporate_actions"))[["stock_id", "date", "dividend", "split_factor"]]
    actions["stock_id"] = actions["stock_id"].astype(str)
    splits = actions[actions["split_factor"] > 0]
    paid = actions[actions["dividend"] > 0].copy()
    # restate each payment in today's share count: divide by every later split
    factor = np.ones(len(paid))
    for stock, events in splits.groupby("stock_id"):
        rows = np.where(paid["stock_id"].values == stock)[0]
        for when, ratio in zip(events["date"].values, events["split_factor"].values):
            later = paid["date"].values[rows] < when
            factor[rows[later]] *= ratio
    paid["per_share"] = paid["dividend"].values / factor
    # a special: more than 2.5x the median of the stock's previous 8 payments
    paid = paid.sort_values(["stock_id", "date"])
    typical = paid.groupby("stock_id")["per_share"].transform(lambda s: s.shift(1).rolling(8, min_periods=3).median())
    paid = paid[~(paid["per_share"] > 2.5 * typical)]
    by_day = paid.pivot_table(index="date", columns="stock_id", values="per_share", aggfunc="sum")
    days = prices().index
    by_day = by_day.reindex(days.union(by_day.index)).fillna(0)
    return by_day.rolling("365D").sum().reindex(days).reindex(columns=prices().columns).fillna(0)


def dividend_universe(top=500):
    """The `top` largest US companies."""
    return largest(company_ids(us_only=True), top)


def dividend_growth(years=3, top=500):
    """Trailing dividend per share higher than a year earlier, `years` years in a row."""
    paid = dividends_12m()
    year = 252
    raised = paid > 0
    for k in range(years):
        raised &= paid.shift(k * year) > paid.shift((k + 1) * year)
    raised &= paid.shift(years * year) > 0
    return as_mask(raised) & dividend_universe(top)


def equal_weight(universe="sp500"):
    if universe == "sp500":
        return as_mask(get("us_index_constituents:sp500"))
    return largest(company_ids(), int(universe.replace("top", "")))


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


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()
    monthly = curve.resample("ME").last().pct_change().dropna()
    downside = math.sqrt((monthly.clip(upper=0) ** 2).mean())
    return {
        "cagr": curve.iloc[-1] ** (1 / years) - 1,
        "sharpe": daily.mean() / daily.std() * math.sqrt(252),  # daily, no risk-free rate
        "monthly_sortino": monthly.mean() / downside * math.sqrt(12),
        "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 compare(curve, benchmark):
    """Daily Sharpe of the strategy vs QQQ: full window and each half."""
    days = curve.index
    mid = days[days.searchsorted(days[0] + (days[-1] - days[0]) / 2)]
    out = {"mid": str(mid.date())}
    for label, part in (("full", slice(None)), ("first_half", slice(None, mid)), ("second_half", slice(mid, None))):
        s, b = stats(curve.loc[part]), stats(benchmark.loc[part])
        out[label] = {"strategy": s, "qqq": b, "beats_qqq": bool(s["sharpe"] > b["sharpe"])}
    return out


def holdings(position, start=START, resample="M"):
    """Average holdings, share of rebalances spent in cash, one-way turnover per year."""
    rule, per_year = {"W": ("W-FRI", 52), "Q": ("QE", 4), "Y": ("YE", 1)}.get(resample, ("ME", 12))
    held = position.loc[start:].astype(float).resample(rule).last()
    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() * per_year,
    }


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, start=START):
    sector = get("us_company_profile").set_index("stock_id")["sector"]
    held = position.loc[start:].astype(float).resample("ME").last()
    weights = held.div(held.sum(axis=1), axis=0)
    mix = weights.T.groupby(sector.reindex(weights.columns).fillna("Unknown").values).sum().T.mean()
    return mix[mix > 0].sort_values(ascending=False).to_dict()


# ---------------------------------------------------------------------------
# 5. Run
# ---------------------------------------------------------------------------
SCREENS = {
    "rule_of_40": dict(build=rule_of_40, resample="M", base={"threshold": 0.40, "min_cap": 2e9},
                       grid={"threshold": [0.30, 0.40, 0.50], "min_cap": [1e9, 2e9, 5e9]}),
    "magic_formula": dict(build=magic_formula, resample="Q", base={"n": 30, "top": 1000},
                          grid={"n": [20, 30, 50], "top": [500, 1000, 2000]}),
    "canslim": dict(build=canslim, resample="W", base={"growth": 0.25, "rs_min": 0.80, "near_high": 0.85},
                    grid={"growth": [0.20, 0.25, 0.30], "rs_min": [0.70, 0.80, 0.90], "near_high": [0.80, 0.85, 0.90]}),
    "graham_number": dict(build=graham_number, resample="M", base={"multiple": 22.5, "top": 1000},
                          grid={"multiple": [15, 22.5, 30], "top": [500, 1000, 1500]}),
    "dividend_growth": dict(build=dividend_growth, resample="M", start=DIVIDEND_START, base={"years": 3, "top": 500},
                            grid={"years": [2, 3, 4], "top": [300, 500, 1000]}),
    "equal_weight_sp500": dict(build=equal_weight, resample="Q", base={"universe": "sp500"},
                               grid={"resample": ["M", "Q", "Y"], "universe": ["sp500", "top500", "top1000"]}),
}


def run_screen(spec, params, **costs):
    params = dict(params)
    resample = params.pop("resample", spec["resample"])
    position = spec["build"](**params)
    start = spec.get("start", START)
    return position, backtest(position, start=start, resample=resample, **costs)


def summary_line(name, r):
    f = r["full"]
    return (f"{name:20s} Sharpe {f['strategy']['sharpe']:.2f} vs QQQ {f['qqq']['sharpe']:.2f} | halves "
            f"{r['first_half']['strategy']['sharpe']:.2f}/{r['first_half']['qqq']['sharpe']:.2f}, "
            f"{r['second_half']['strategy']['sharpe']:.2f}/{r['second_half']['qqq']['sharpe']:.2f} | "
            f"CAGR {f['strategy']['cagr']:.1%} MDD {f['strategy']['max_drawdown']:.1%}")


def grid_results(spec, qqq):
    rows = []
    for combo in itertools.product(*spec["grid"].values()):
        params = dict(zip(spec["grid"], combo))
        _, curve = run_screen(spec, params)
        r = compare(curve, qqq)
        rows.append({"params": params, "sharpe": r["full"]["strategy"]["sharpe"], "qqq": r["full"]["qqq"]["sharpe"],
                     "first_half": r["first_half"]["strategy"]["sharpe"], "second_half": r["second_half"]["strategy"]["sharpe"],
                     "cagr": r["full"]["strategy"]["cagr"], "max_drawdown": r["full"]["strategy"]["max_drawdown"]})
    return {"rows": rows, "wins": sum(row["sharpe"] > row["qqq"] for row in rows), "total": len(rows)}


def extras(qqq_long, curves):
    """Universe controls, rule variants, and the dividend start-date and streak tests."""
    out = {}
    days = qqq_long.index
    *_, market_up, canslim_universe = canslim_inputs()
    controls = {
        "rule_of_40": rule_of_40_universe(),
        "magic_formula": magic_formula_universe(),
        "canslim": canslim_universe,
        "graham_number": graham_universe(),
    }
    out["universe_equal_weight"] = {k: stats(backtest(u, resample=SCREENS[k]["resample"])) for k, u in controls.items()}
    magic = magic_formula()
    out["variants"] = {
        "canslim_stop_8pct": compare(backtest(canslim(), resample="W", stop_loss=0.08), qqq_long),
        "canslim_no_market_filter": compare(backtest(canslim(market_filter=False), resample="W"), qqq_long),
        "canslim_market_up_share": float(market_up.loc[START:].mean()),
        "magic_formula_monthly": compare(backtest(magic, resample="M"), qqq_long),
        "magic_formula_yearly": compare(backtest(magic, resample="Y"), qqq_long),
    }
    ew = curves["equal_weight_sp500"]
    out["equal_weight_vs_rsp_daily_corr"] = float(ew.pct_change().corr(fund("RSP", days).pct_change()))
    out["etfs"] = {t: stats(fund(t, days)) for t in ETFS}
    div_days = prices().loc[DIVIDEND_START:].index
    out["etfs_dividend_window"] = {t: stats(fund(t, div_days)) for t in ETFS}

    # dividend growth: the same rule started on different dates
    rule = dividend_growth(3, 500)
    starts = {}
    for start in ["2020-01-02", "2020-07-01", "2021-01-04", "2021-07-01", "2022-01-03", "2022-07-01", "2023-01-03"]:
        curve = backtest(rule, start=start)
        span = curve.index
        starts[start] = {"strategy": stats(curve), "qqq": stats(fund("QQQ", span)), "spy": stats(fund("SPY", span))}
    out["dividend_start_dates"] = starts
    out["dividend_shorter_records"] = {
        f"{years}y_from_{start[:4]}": compare(backtest(dividend_growth(years, 500), start=start), fund("QQQ", prices().loc[start:].index))
        for years, start in ((1, "2018-01-02"), (2, "2019-01-02"))
    }
    # dividend growth: what each dividend record earned inside the same 500 companies
    paid = dividends_12m()
    universe = dividend_universe(500)
    raised = [paid.shift(k * 252) > paid.shift((k + 1) * 252) for k in range(4)]
    streak, run = raised[0].astype(int), raised[0]
    for r in raised[1:]:
        run = run & r
        streak = streak + run.astype(int)
    payer = paid > 0
    cut = payer & (paid < paid.shift(252))
    groups = {
        "no_dividend": ~payer, "cut": cut, "unchanged": payer & ~cut & (streak == 0),
        "raised_1y": payer & (streak == 1), "raised_2y": payer & (streak == 2),
        "raised_3y": payer & (streak == 3), "raised_4y_plus": payer & (streak >= 4), "all_500": payer | ~payer,
    }
    out["dividend_record_groups"] = {}
    for k, g in groups.items():
        position = as_mask(g) & universe
        out["dividend_record_groups"][k] = {**stats(backtest(position, start=DIVIDEND_START)), **holdings(position, DIVIDEND_START)}
    out["dividend_growth_yearly"] = {"strategy": yearly(curves["dividend_growth"]), "QQQ": yearly(fund("QQQ", div_days)),
                                     "SPY": yearly(fund("SPY", div_days)), "VIG": yearly(fund("VIG", div_days))}
    out["dividend_growth_sector"] = sector_mix(rule, DIVIDEND_START)
    yield_now = paid / get("us_price:close").reindex_like(paid)
    out["dividend_growth_median_yield"] = float(yield_now[rule].loc[DIVIDEND_START:].resample("ME").last().stack().median())
    return out


def main():
    qqq_long = fund("QQQ", prices().loc[START:].index)
    curves, results = {}, {}
    for name, spec in SCREENS.items():
        start = spec.get("start", START)
        position, curve = run_screen(spec, spec["base"])
        qqq = fund("QQQ", curve.index)
        with_cost = backtest(position, start=start, resample=spec["resample"], fee_ratio=0.001, tax_ratio=0)
        results[name] = {"params": spec["base"], **compare(curve, qqq), "cost_0.1pct": compare(with_cost, qqq),
                         **holdings(position, start, spec["resample"]), "sector": sector_mix(position, start)}
        if start == START:
            results[name]["yearly"] = yearly(curve)
        curves[name] = curve
        print(summary_line(name, results[name]), flush=True)

    days = qqq_long.index
    metrics = {"window": [str(days[0].date()), str(days[-1].date())], "dividend_window_start": DIVIDEND_START,
               "screens": results, "qqq": stats(qqq_long), "spy": stats(fund("SPY", days)),
               "qqq_yearly": yearly(qqq_long), "spy_yearly": yearly(fund("SPY", days))}
    if "--grid" in sys.argv:
        metrics["grid"] = {}
        for name, spec in SCREENS.items():
            qqq = fund("QQQ", curves[name].index)
            metrics["grid"][name] = grid_results(spec, qqq)
            print(f"{name:20s} grid: {metrics['grid'][name]['wins']}/{metrics['grid'][name]['total']} settings beat QQQ", flush=True)
    if "--extras" in sys.argv:
        metrics["extras"] = extras(qqq_long, curves)

    frame = pd.DataFrame(curves)
    for t in ["QQQ", "SPY", "RSP"]:
        frame[t] = fund(t, frame.index)
    frame.round(6).to_csv("data.csv")
    with open("metrics.json", "w") as f:
        json.dump(metrics, f, indent=1, default=float)


if __name__ == "__main__":
    main()