"""
Ten academic stock factors on US equities, long-only, versus QQQ (total return).

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

What it does
  1. Builds a point-in-time liquid US universe: common stocks (delisted names included),
     ex financials, top 1,000 by 60-day dollar volume, price > $5, > $5M/day.
  2. Scores every stock on ten published factors (one function each, below).
  3. Holds the top 30 stocks of each factor, equal weight, rebalanced monthly,
     with finlab's default US costs (USMarket: commission 0, no transfer tax).
  4. Compares each portfolio with QQQ buy-and-hold over the same days:
     full window, first half, second half.
  5. Optional: the 3 x 3 parameter grid and the decile / rank-IC test.

Run:  python strategy.py            (10 base portfolios, a few minutes)
      python strategy.py --grid     (adds the 90-portfolio parameter grid)
      python strategy.py --deciles  (adds the decile and rank-IC test)
"""
import json
import math
import sys

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 a year of prices
BASE_UNIVERSE = 1000      # liquid names kept each day
BASE_HOLDINGS = 30        # stocks held per factor

data.set_market("us")


# ---------------------------------------------------------------------------
# 1. Universe
# ---------------------------------------------------------------------------
def common_stock_ids():
    """Common stocks from the company profile, including delisted ones, ex financials.

    etl:us_common_stock_filter only lists companies that still trade today, which
    would drop every bankruptcy and takeover from the test (survivorship bias).
    """
    profile = data.get("us_company_profile")
    is_common = (
        (profile["security_type"] == "common_stock")
        & ~profile["is_etf"].fillna(False).astype(bool)
        & ~profile["is_fund"].fillna(False).astype(bool)
        & ~profile["is_adr"].fillna(False).astype(bool)
        & (profile["sector"] != "Financial Services")
    )
    return sorted(profile.loc[is_common, "stock_id"].astype(str))


def liquid_universe(top_n=BASE_UNIVERSE):
    close = data.get("us_price:close")
    volume = data.get("us_price:volume")
    ids = close.columns.intersection(common_stock_ids())
    close, volume = close[ids], volume[ids]

    dollar_volume = (close * volume).rolling(60, min_periods=20).mean()

    # Drop a stock for a year after any +/-50% day (split errors, ticker reuse).
    # The window only looks backwards, so no future information leaks in.
    daily_return = adj_close()[ids].pct_change(fill_method=None)
    extreme_day = (daily_return.abs() >= 0.5).astype(float)
    recently_extreme = extreme_day.rolling(252, min_periods=1).max().shift(1).fillna(0) > 0

    return (
        dollar_volume.is_largest(top_n)
        & (dollar_volume > 5e6)
        & (close > 5)
        & ~recently_extreme
    )


# ---------------------------------------------------------------------------
# 2. Data helpers
# ---------------------------------------------------------------------------
_cache = {}


def adj_close():
    if "adj" not in _cache:
        _cache["adj"] = data.get("us_price:adj_close")
    return _cache["adj"]


def quarterly(key):
    """A quarterly statement item on a shared quarter x ticker grid.

    Values are dated by SEC filing date (key_date), so they become usable
    only after the 10-Q / 10-K is public. No extra shift is needed.
    """
    if key not in _cache:
        quarters = data.get("us_balance_sheet:total_assets").index
        _cache[key] = data.get(key).reindex(index=quarters, columns=adj_close().columns)
    return _cache[key]


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


def market_cap():
    return data.get("etl:us_market_value")


def average_assets():
    assets = quarterly("us_balance_sheet:total_assets")
    return (assets + assets.shift(4)) / 2


def market_return():
    spy = data.get("us_fund_price:adj_close")["SPY"]
    return spy.reindex(adj_close().index).ffill().pct_change(fill_method=None)


# ---------------------------------------------------------------------------
# 3. Ten factors (higher score = buy)
# ---------------------------------------------------------------------------
def value_book_to_market():
    """Fama & French (1992): book equity / market cap."""
    equity = quarterly("us_balance_sheet:total_stockholders_equity")
    bm = equity / market_cap()
    return bm[bm > 0]


def gross_profitability():
    """Novy-Marx (2013): trailing gross profit / total assets."""
    assets = quarterly("us_balance_sheet:total_assets")
    return trailing_12m("us_income_statement:gross_profit") / assets[assets > 0]


def low_asset_growth():
    """Cooper, Gulen & Schill (2008): year-over-year total asset growth, low is good."""
    assets = quarterly("us_balance_sheet:total_assets")
    previous_assets = assets.shift(4)
    growth = assets / previous_assets - 1
    return -growth[previous_assets > 0]


def low_accruals():
    """Sloan (1996), cash-flow version: (net income - operating cash flow) / avg assets."""
    accruals = (
        trailing_12m("us_income_statement:net_income")
        - trailing_12m("us_cash_flow:operating_cash_flow")
    ) / average_assets()
    return -accruals


def net_buyback_yield():
    """Pontiff & Woodgate (2008), Ikenberry et al. (1995): -(net share issuance) / market cap."""
    issuance = trailing_12m("us_cash_flow:net_common_stock_issuance")
    cap = market_cap()
    return -(issuance / cap[cap > 0])


def near_52_week_high(window=252):
    """George & Hwang (2004): price / highest price of the past year."""
    price = adj_close()
    return price / price.rolling(window, min_periods=int(window * 0.8)).max()


def _rolling_market_model(window, min_obs):
    """Daily stock returns, market variance and betas for the same rolling window."""
    stock = adj_close().pct_change(fill_method=None)
    market = market_return()
    market_mean = market.rolling(window, min_periods=min_obs).mean()
    market_var = market.rolling(window, min_periods=min_obs).var()
    stock_mean = stock.rolling(window, min_periods=min_obs).mean()
    cross_mean = stock.mul(market, axis=0).rolling(window, min_periods=min_obs).mean()
    covariance = (cross_mean - stock_mean.mul(market_mean, axis=0)) * (window / (window - 1))
    return stock, market_var, covariance.div(market_var, axis=0)


def rolling_beta(window=252):
    _, _, beta = _rolling_market_model(window, max(int(window * 0.8), 20))
    return beta


def low_beta(window=252):
    """Frazzini & Pedersen (2014), long leg only: low market beta vs SPY."""
    return -rolling_beta(window)


def low_idiosyncratic_vol(window=21):
    """Ang, Hodrick, Xing & Zhang (2006): volatility of CAPM residuals, low is good."""
    min_obs = max(int(window * 0.8), min(15, window))
    stock, market_var, beta = _rolling_market_model(window, min_obs)
    residual_var = stock.rolling(window, min_periods=min_obs).var() - beta.pow(2).mul(market_var, axis=0)
    return -np.sqrt(residual_var.clip(lower=0))


def piotroski_f_score():
    """Piotroski (2000): nine pass/fail accounting signals, 0-9."""
    net_income = trailing_12m("us_income_statement:net_income")
    cfo = trailing_12m("us_cash_flow:operating_cash_flow")
    assets = quarterly("us_balance_sheet:total_assets")
    begin_assets = assets.shift(4)
    begin_assets = begin_assets[begin_assets > 0]
    roa = net_income / begin_assets
    cfo_to_assets = cfo / begin_assets
    leverage = quarterly("us_balance_sheet:long_term_debt").fillna(0) / average_assets()
    current_ratio = (
        quarterly("us_balance_sheet:total_current_assets")
        / quarterly("us_balance_sheet:total_current_liabilities")
    )
    issuance = trailing_12m("us_cash_flow:net_common_stock_issuance")
    revenue = trailing_12m("us_income_statement:revenue")
    gross_margin = trailing_12m("us_income_statement:gross_profit") / revenue[revenue > 0]
    turnover = revenue / begin_assets

    signals = [
        roa > 0,
        cfo > 0,
        roa > roa.shift(4),
        cfo_to_assets > roa,
        leverage < leverage.shift(4),
        current_ratio > current_ratio.shift(4),
        issuance <= 0,
        gross_margin > gross_margin.shift(4),
        turnover > turnover.shift(4),
    ]
    score = sum(s.astype(float) for s in signals)
    return score[roa.notna() & cfo.notna() & roa.shift(4).notna()]


def rank_in(score, universe):
    """Cross-sectional percentile rank among today's universe members."""
    return (score * 1.0)[universe].rank(axis=1, pct=True).reindex(universe.index)


def mean_rank(ranks):
    """Average of available ranks, re-ranked (a stock needs at least one input)."""
    total = ranks[0].fillna(0)
    count = ranks[0].notna().astype(float)
    for r in ranks[1:]:
        total = total + r.fillna(0)
        count = count + r.notna().astype(float)
    return (total / count[count > 0]).rank(axis=1, pct=True)


def quality_minus_junk(universe):
    """Asness, Frazzini & Pedersen (2019): profitability + growth + safety + payout.

    Growth uses 1-year changes (the paper uses 5 years; the data starts in 2016).
    """
    ni = trailing_12m("us_income_statement:net_income")
    cfo = trailing_12m("us_cash_flow:operating_cash_flow")
    gp = trailing_12m("us_income_statement:gross_profit")
    revenue = trailing_12m("us_income_statement:revenue")
    equity = quarterly("us_balance_sheet:total_stockholders_equity")
    assets = quarterly("us_balance_sheet:total_assets")
    avg_assets = average_assets()

    gpoa = gp / assets[assets > 0]
    roe = ni / equity[equity > 0]
    roa = ni / avg_assets[avg_assets > 0]
    cfoa = cfo / avg_assets[avg_assets > 0]
    gmar = gp / revenue[revenue > 0]
    low_acc = -(ni - cfo) / avg_assets[avg_assets > 0]

    profitability = [gpoa, roe, roa, cfoa, gmar, low_acc]
    growth = [x - x.shift(4) for x in (gpoa, roe, roa, cfoa, gmar)]
    leverage = quarterly("us_balance_sheet:total_debt").fillna(0) / assets[assets > 0]
    quarterly_roe = quarterly("us_income_statement:net_income") / equity[equity > 0]
    safety = [-rolling_beta(252), -leverage, -quarterly_roe.rolling(8, min_periods=6).std()]

    groups = [mean_rank([rank_in(x, universe) for x in parts]) for parts in (profitability, growth, safety)]
    groups.append(rank_in(net_buyback_yield(), universe))
    return mean_rank(groups)


# ---------------------------------------------------------------------------
# 4. Portfolios, backtest, statistics
# ---------------------------------------------------------------------------
def top_n(score, universe, n=BASE_HOLDINGS):
    position = score[universe].is_largest(n)
    return position.reindex(universe.index).fillna(False).astype(bool)


def f_score_portfolio(universe, value_share=0.3, min_score=8):
    """Piotroski's rule: high F-score among the cheapest stocks by book-to-market."""
    cheap = rank_in(value_book_to_market(), universe) >= 1 - value_share
    position = (cheap & (piotroski_f_score() >= min_score))[universe]
    return position.reindex(universe.index).fillna(False).astype(bool)


def backtest(position):
    report = sim(position.loc[START:], resample="M", upload=False)
    curve = report.creturn
    # A portfolio that finds no stock in its first months sits in cash until it does.
    days = adj_close().loc[START:].index
    return curve.reindex(days).ffill().fillna(1.0)


def annual_turnover(position):
    """One-way turnover per year from month-end equal weights (price drift ignored)."""
    held = position.loc[START:].resample("ME").last().astype(float)
    counts = held.sum(axis=1)
    weights = held.div(counts.where(counts > 0), axis=0).fillna(0)
    return float((weights.diff().abs().sum(axis=1) / 2).iloc[1:].mean() * 12)


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()
    daily_std = daily.std()
    monthly = curve.resample("ME").last().pct_change().dropna()
    monthly_down = 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() / monthly_down * math.sqrt(12),
        "max_drawdown": (curve / curve.cummax() - 1).min(),
        "volatility": daily_std * math.sqrt(252),
    }


def qqq_curve(days):
    qqq = data.get("us_fund_price:adj_close")["QQQ"]
    qqq = qqq.reindex(days.union(qqq.index)).ffill().reindex(days)
    return qqq / qqq.iloc[0]


def compare(curve, benchmark):
    """Sharpe of strategy vs QQQ over the full window and each half."""
    mid = curve.index[curve.index.searchsorted(curve.index[0] + (curve.index[-1] - curve.index[0]) / 2)]
    out = {}
    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": s["sharpe"] > b["sharpe"]}
    return out


# ---------------------------------------------------------------------------
# 5. Run
# ---------------------------------------------------------------------------
def base_portfolios(universe):
    return {
        "value_bm": top_n(value_book_to_market(), universe),
        "gross_profitability": top_n(gross_profitability(), universe),
        "asset_growth": top_n(low_asset_growth(), universe),
        "accruals": top_n(low_accruals(), universe),
        "buyback_yield": top_n(net_buyback_yield(), universe),
        "high_52w": top_n(near_52_week_high(252), universe),
        "low_beta": top_n(low_beta(252), universe),
        "low_ivol": top_n(low_idiosyncratic_vol(21), universe),
        "fscore": f_score_portfolio(universe, 0.3, 8),
        "qmj": top_n(quality_minus_junk(universe), universe),
    }


def parameter_grid():
    """3 x 3 neighbours around each base setting (base in the middle)."""
    holdings = [20, 30, 50]
    sizes = [500, 1000, 1500]
    universes = {size: liquid_universe(size) for size in sizes}
    u = universes[BASE_UNIVERSE]
    grid = {}
    fundamentals = {
        "value_bm": value_book_to_market, "gross_profitability": gross_profitability,
        "asset_growth": low_asset_growth, "accruals": low_accruals, "buyback_yield": net_buyback_yield,
    }
    for name, fn in fundamentals.items():
        score = fn()
        grid[name] = {(size, n): top_n(score, universes[size], n) for size in sizes for n in holdings}
    grid["qmj"] = {}
    for size in sizes:
        score = quality_minus_junk(universes[size])
        for n in holdings:
            grid["qmj"][(size, n)] = top_n(score, universes[size], n)
    lookbacks = {
        "high_52w": (near_52_week_high, [189, 252, 315]),
        "low_beta": (low_beta, [126, 252, 378]),
        "low_ivol": (low_idiosyncratic_vol, [10, 21, 42]),
    }
    for name, (fn, windows) in lookbacks.items():
        grid[name] = {}
        for w in windows:
            score = fn(w)
            for n in holdings:
                grid[name][(w, n)] = top_n(score, u, n)
    grid["fscore"] = {
        (share, k): f_score_portfolio(u, share, k) for share in [0.2, 0.3, 0.4] for k in [7, 8, 9]
    }
    return grid


def decile_test(universe, scores):
    """Equal-weight decile returns, month-end signal, next-day entry (same timing as sim)."""
    price = adj_close()
    days = price.loc[START:].index
    month_ends = pd.DatetimeIndex(pd.Series(days, index=days).groupby([days.year, days.month]).last().values)
    entry_pos = days.searchsorted(month_ends) + 1
    keep = entry_pos < len(days)
    month_ends, entries = month_ends[keep], days[entry_pos[keep]]
    exits = list(entries[1:]) + [days[-1]]
    px = price.reindex(columns=universe.columns)
    forward = pd.DataFrame({m: px.loc[x] / px.loc[e] - 1 for m, e, x in zip(month_ends, entries, exits)}).T
    result = {}
    for name, score in scores.items():
        s = (score * 1.0)[universe].reindex(index=universe.index).reindex(month_ends)
        decile = np.ceil(s.rank(axis=1, pct=True) * 10).clip(1, 10)
        dec_ret = pd.DataFrame({d: forward.where(decile == d).mean(axis=1) for d in range(1, 11)})
        spread = dec_ret[10] - dec_ret[1]
        ic = pd.Series([s.loc[m].corr(forward.loc[m], method="spearman") for m in month_ends], index=month_ends)
        half = len(spread) // 2
        result[name] = {
            "decile_cagr": {int(d): float((1 + dec_ret[d]).prod() ** (12 / len(dec_ret)) - 1) for d in dec_ret},
            "spread_annual": float(spread.mean() * 12),
            "spread_t": float(spread.mean() / spread.std() * math.sqrt(len(spread))),
            "spread_first_half": float(spread.iloc[:half].mean() * 12),
            "spread_second_half": float(spread.iloc[half:].mean() * 12),
            "rank_ic": float(ic.mean()),
            "rank_ic_t": float(ic.mean() / ic.std() * math.sqrt(ic.notna().sum())),
        }
    return result


def main():
    universe = liquid_universe()
    portfolios = base_portfolios(universe)
    curves = {name: backtest(pos) for name, pos in portfolios.items()}
    days = next(iter(curves.values())).index
    qqq = qqq_curve(days)
    spy = data.get("us_fund_price:adj_close")["SPY"].reindex(days).ffill()
    curves["equal_weight_universe"] = backtest(universe)

    results = {name: compare(c, qqq) for name, c in curves.items()}
    for name, pos in portfolios.items():
        results[name]["annual_turnover"] = annual_turnover(pos)
        results[name]["avg_holdings"] = float(pos.loc[START:].resample("ME").last().sum(axis=1).mean())
    for name, r in results.items():
        f = r["full"]
        print(f"{name:22s} Sharpe {f['strategy']['sharpe']:.2f} vs QQQ {f['qqq']['sharpe']:.2f} | "
              f"halves {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%}", flush=True)

    metrics = {"window": [str(days[0].date()), str(days[-1].date())], "base": results,
               "qqq": stats(qqq), "spy": stats(spy / spy.iloc[0])}

    if "--grid" in sys.argv:
        metrics["grid"] = {}
        for name, combos in parameter_grid().items():
            rows = []
            for params, pos in combos.items():
                r = compare(backtest(pos), qqq)
                rows.append({"params": list(params), **{k: {"sharpe": v["strategy"]["sharpe"], "qqq": v["qqq"]["sharpe"]} for k, v in r.items()},
                             "cagr": r["full"]["strategy"]["cagr"], "max_drawdown": r["full"]["strategy"]["max_drawdown"],
                             "monthly_sortino": r["full"]["strategy"]["monthly_sortino"],
                             "avg_holdings": float(pos.loc[START:].resample("ME").last().sum(axis=1).mean())})
            wins = sum(row["full"]["sharpe"] > row["full"]["qqq"] for row in rows)
            metrics["grid"][name] = {"combos": rows, "wins": wins}
            print(f"{name:22s} grid: {wins}/9 settings beat QQQ on Sharpe", flush=True)

    if "--deciles" in sys.argv:
        scores = {
            "value_bm": value_book_to_market(), "gross_profitability": gross_profitability(),
            "asset_growth": low_asset_growth(), "accruals": low_accruals(), "buyback_yield": net_buyback_yield(),
            "high_52w": near_52_week_high(252), "low_beta": low_beta(252), "low_ivol": low_idiosyncratic_vol(21),
            "fscore": piotroski_f_score() + rank_in(value_book_to_market(), universe) * 0.5,
            "qmj": quality_minus_junk(universe),
        }
        metrics["deciles"] = decile_test(universe, scores)

    pd.DataFrame({**curves, "QQQ": qqq, "SPY": spy / spy.iloc[0]}).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()