"""
Piotroski F-Score rules on S&P 500 stocks, 2018 to 2026 (FinLab research, 2026-11).

The score adds nine accounting checks, as in Piotroski (2000), on trailing four-quarter figures
against a year earlier. Every member is scored; financials stay in. Designs:

  fscore_7             every stock scoring at least f
                       grid: f 6/7/8 x weekly/monthly/quarterly, original 7, monthly
  fscore_8             the same around 8
                       grid: f 7/8/9 x weekly/monthly/quarterly, original 8, monthly
  fscore_cheap_growth  the cheapest x of members by price to book, scoring at least k, with the
                       latest quarter's revenue above a year earlier, monthly
                       grid: x 20/30/40% x k 6/7/8, original 30%, 7

Each holds equal weights with at most 5% in a stock. All start in 2018-05, the first month at
least 70% of members have a score: statements for S&P 500 companies start in 2016-Q1.

Signals are read at a close and traded at the next close. Quarterly figures count from their
filing date. The universe is each day's S&P 500 members, one share class per company.

"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 70% of its 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 still beats QQQ without the calendar year of its largest lead,
  5. if it averages more than 10% cash, it beats QQQ on returns above T-bills (BIL),
  6. a combined design beats its single-signal base over the same dates.

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

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

END = "2026-10-30"
COST_PER_SIDE = 0.001
MAX_WEIGHT = 0.05
OUT = Path(__file__).resolve().parent
FUNDS = ("QQQ", "SPY", "RSP", "SPLV", "USMV", "BIL")

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


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


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


@cache
def funds():
    """Fund prices with dividends on stock-market trading days (us_fund_price has holiday rows)."""
    return get("us_fund_price:adj_close")[list(FUNDS)].reindex(prices().index)


@cache
def profile():
    return pd.DataFrame(get("us_company_profile")).astype({"stock_id": str}).drop_duplicates("stock_id").set_index("stock_id")


@cache
def members():
    """Trading day x ticker: True while the stock is in the S&P 500 (one share class per company)."""
    m = get("us_index_constituents:sp500")
    names = profile()["company_name"].reindex(sorted(m.columns.intersection(prices().columns)))
    m = m[names.index].astype(bool).reindex(prices().index, 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   # renamed tickers keep their own dates
    return FinlabDataFrame(m)


@cache
def sectors():
    """Today's sector labels. The table keeps no history of them."""
    return profile()["sector"].reindex(members().columns)

RHYTHMS = ("W", "M", "Q")

DESIGNS = {  # name: (grid, original setting)
    "fscore_7": (((6, 7, 8), RHYTHMS), (7, "M")),
    "fscore_8": (((7, 8, 9), RHYTHMS), (8, "M")),
    "fscore_cheap_growth": (((0.2, 0.3, 0.4), (6, 7, 8)), (0.3, 7)),
}
COMBINES = {"fscore_cheap_growth": "fscore_7"}


@cache
def quarterly(key):
    """A statement item on a shared quarter x ticker grid."""
    quarters = get("us_balance_sheet:total_assets").index
    return get(key).reindex(index=quarters, columns=members().columns)


def on_days(frame, unknown):
    """A quarterly frame on the trading-day grid, from each filing date. A blank quarter becomes
    `unknown` first: left blank, the filing would keep showing the quarter before."""
    return frame.fillna(unknown).index_str_to_date().reindex(index=prices().index, method="ffill")


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


@cache
def market_value():
    return get("etl:us_market_value").reindex(index=prices().index, columns=members().columns)


# ---------------------------------------------------------------------------
# Piotroski F-Score
# ---------------------------------------------------------------------------
@cache
def f_score():
    """Piotroski (2000), the nine signals as in /blog/us-academic-factors-backtest."""
    income = trailing("us_income_statement:net_income")
    cash_flow = trailing("us_cash_flow:operating_cash_flow")
    assets = quarterly("us_balance_sheet:total_assets")
    begin = assets.shift(4)
    begin = begin[begin > 0]
    roa = income / begin
    cash_to_assets = cash_flow / begin
    leverage = quarterly("us_balance_sheet:long_term_debt").fillna(0) / ((assets + assets.shift(4)) / 2)
    current = quarterly("us_balance_sheet:total_current_assets") / quarterly("us_balance_sheet:total_current_liabilities")
    issued = trailing("us_cash_flow:net_common_stock_issuance")
    revenue = trailing("us_income_statement:revenue")
    margin = trailing("us_income_statement:gross_profit") / revenue[revenue > 0]
    turnover = revenue / begin
    signals = [
        roa > 0, cash_flow > 0, roa > roa.shift(4), cash_to_assets > roa,
        leverage < leverage.shift(4), current > current.shift(4), issued <= 0,
        margin > margin.shift(4), turnover > turnover.shift(4),
    ]
    score = sum(s.astype(float) for s in signals)
    return on_days(score[roa.notna() & cash_flow.notna() & roa.shift(4).notna()], -1)


@cache
def price_to_book():
    equity = on_days(quarterly("us_balance_sheet:total_stockholders_equity"), -1)
    return market_value() / equity[equity > 0]


@cache
def revenue_up():
    revenue = quarterly("us_income_statement:revenue")
    earlier = revenue.shift(4)
    return on_days(revenue[revenue > 0] / earlier[earlier > 0] - 1, -1) > 0


@cache
def position(name, *setting):
    if name == "fscore_cheap_growth":
        share, floor = setting
        cheap = price_to_book()[members()].rank(axis=1, pct=True) <= share
        return cheap & (f_score() >= floor) & revenue_up()
    return (f_score() >= setting[0]) & members()


@cache
def backtest(name, setting, fee_ratio=None):
    table = position(name, *setting).loc["2018-05-01":]
    return simulate(table, "M" if name == "fscore_cheap_growth" else setting[1], MAX_WEIGHT, fee_ratio)


# ---------------------------------------------------------------------------
# Backtest and statistics
# ---------------------------------------------------------------------------
def simulate(table, rhythm, limit, fee_ratio=None):
    table = table.loc[:, table.any()]
    report = sim(table, resample=rhythm, trade_at_price="close", fee_ratio=fee_ratio,
                 position_limit=limit, market="us", upload=False)
    curve = report.creturn
    moved = ((curve / curve.iloc[0] - 1).abs() > 1e-9).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):], report


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


def bench(ticker, index):
    return funds()[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 cagr(curve):
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return (curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1


def stats(curve):
    out = {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date())}
    lines = (("", curve), *((f"{t.lower()}_", bench(t, curve.index)) for t in FUNDS[:-1]))
    for label, series in lines:
        out[f"{label}sharpe"] = round(daily_sharpe(series.pct_change().dropna()), 6)
        out[f"{label}cagr"] = round(cagr(series), 6)
        out[f"{label}mdd"] = round((series / series.cummax() - 1).min(), 6)
    out["volatility"] = round(curve.pct_change().std() * 252 ** 0.5, 6)
    out["qqq_volatility"] = round(bench("QQQ", curve.index).pct_change().std() * 252 ** 0.5, 6)
    out["excess_sharpe"] = round(excess_sharpe(curve), 6)
    out["qqq_excess_sharpe"] = round(excess_sharpe(bench("QQQ", curve.index)), 6)
    out["beats_qqq"] = bool(out["sharpe"] > out["qqq_sharpe"])
    return out


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 halves(curve):
    middle = curve.index[len(curve) // 2]
    return [stats(part) for part in (curve.loc[:middle], curve.loc[middle:])]


def without_best_year(curve):
    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 study(name):
    axes, original = DESIGNS[name]
    grid = {str(s): stats(backtest(name, s)[0]) for s in product(*axes)}
    curve, report = backtest(name, original)
    first, second = halves(curve)
    weights = report.position.fillna(0).astype(float)
    weights = weights.loc[curve.index[0]:]
    held = (weights != 0).sum(axis=1)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    sector_share = weights.abs().T.groupby(sectors().reindex(weights.columns)).sum().T
    result = {
        "name": name, "original": str(original), "original_stats": grid[str(original)], "grid": grid,
        "grid_wins": sum(c["beats_qqq"] for c in grid.values()), "grid_size": len(grid),
        "rsp_wins": int(sum(c["sharpe"] > c["rsp_sharpe"] for c in grid.values())),
        "spy_wins": int(sum(c["sharpe"] > c["spy_sharpe"] for c in grid.values())),
        "first_half": first, "second_half": second,
        "with_cost": stats(backtest(name, original, COST_PER_SIDE)[0].loc[curve.index[0]:]),
        "yearly": yearly(curve), "yearly_qqq": yearly(bench("QQQ", curve.index)), "yearly_rsp": yearly(bench("RSP", curve.index)),
        "without_best_year": without_best_year(curve),
        "holdings": {"mean": round(held.mean(), 2), "median": float(held.median()), "min": int(held.min()), "max": int(held.max())},
        "invested": round(weights.abs().sum(axis=1).clip(upper=1).mean(), 6),
        "turnover": round(weights.diff().abs().sum(axis=1).sum() / 2 / years, 3),
        "sector_mix": {k: round(v, 4) for k, v in sector_share.div(sector_share.sum(axis=1), axis=0).mean().sort_values(ascending=False).items()},
    }
    o = result["original_stats"]
    if name in COMBINES:
        base = backtest(COMBINES[name], DESIGNS[COMBINES[name]][1])[0].loc[curve.index[0]:]
        result["base_same_window"] = stats(base)
    result["passes"] = bool(
        result["grid_wins"] >= 0.7 * len(grid) and first["beats_qqq"] and second["beats_qqq"]
        and o["beats_qqq"] and result["with_cost"]["beats_qqq"]
        and result["without_best_year"]["sharpe"] > result["without_best_year"]["qqq_sharpe"]
        and (result["invested"] >= 0.9 or o["excess_sharpe"] > o["qqq_excess_sharpe"])
        and (name not in COMBINES or o["sharpe"] > result["base_same_window"]["sharpe"]))
    return result


def calibration():
    """Equal weight on every member should track RSP."""
    curve = simulate(members().loc["2017-08-01":], "M", 1)[0]
    paired = pd.concat([curve.pct_change(), bench("RSP", curve.index).pct_change()], axis=1).dropna()
    cell = stats(curve)
    return {"start": cell["start"], "sharpe": cell["sharpe"], "rsp_sharpe": cell["rsp_sharpe"],
            "correlation": round(paired.corr().iloc[0, 1], 4)}


def coverage(day):
    """How many members have a price on a day, before one share class is kept."""
    listed = get("us_index_constituents:sp500").astype(bool).reindex(prices().index, method="ffill").loc[day].fillna(False)
    names = listed.index[listed.astype(bool)]
    return {"day": str(day.date()), "members": len(names), "with_price": int(prices().reindex(columns=names).loc[day].notna().sum()),
            "one_share_class": int(members().loc[day].sum())}


# ---------------------------------------------------------------------------
# Controls added after seeing the results
# ---------------------------------------------------------------------------
def score_bands():
    """Everything monthly and equal weight, so each row is a slice of the S&P 500 by score."""
    score = f_score().reindex(columns=members().columns)
    inside = members().astype(bool)
    rows = {"0-4": (score >= 0) & (score <= 4)}
    rows.update({str(k): score == k for k in (5, 6, 7, 8, 9)})
    rows["no_score"] = score < 0
    out = {}
    for label, picks in rows.items():
        picks = (picks & inside).loc["2018-05-01":]
        curve, _ = simulate(picks, "M", 1)
        held = picks.sum(axis=1).resample("ME").last()
        out[label] = {**stats(curve), "held": round(held.mean(), 1)}
    return out


def financials():
    """Average share of each score band that is financial, month ends."""
    score = f_score().reindex(columns=members().columns)
    inside = members().astype(bool)
    bank = sectors().reindex(score.columns) == "Financial Services"
    ends = score.loc["2018-05-01":].resample("ME").last().index
    days = [score.index[score.index <= d][-1] for d in ends]
    rows = {}
    for label, picks in (("all", score >= 0), ("f7", score >= 7), ("f8", score >= 8), ("f9", score >= 9)):
        picks = (picks & inside).loc[days]
        rows[label] = round((picks.loc[:, bank].sum(axis=1).sum()) / picks.sum(axis=1).sum(), 4)
    return rows


def no_score_snapshot(day="2024-12-31", last="2024-US-Q3", first="2021-US-Q4"):
    """Who has no score on one day: members with filings and blank quarters, not missing companies."""
    score = f_score().reindex(columns=members().columns)
    inside = members().astype(bool)
    names = score.columns[(score.loc[day] < 0) & inside.loc[day]]
    filed = quarterly("us_income_statement:net_income").loc[first:last, names].notna().sum()
    return {"day": day, "members": int(inside.loc[day].sum()), "no_score": len(names), "quarters_checked": 12,
            "with_every_quarter": int((filed == 12).sum()), "with_some_quarters": int(filed.between(1, 11).sum()),
            "with_none": int((filed == 0).sum())}


def extras():
    return {"score_bands": score_bands(), "financial_share": financials(), "no_score_snapshot": no_score_snapshot()}


def main():
    designs = {name: study(name) for name in DESIGNS}
    curves = {name: backtest(name, DESIGNS[name][1])[0] for name in DESIGNS}
    table = pd.DataFrame(curves)
    for ticker in FUNDS[:-1]:
        table[ticker] = bench(ticker, table.index)
    table.div(table.bfill().iloc[0]).round(6).rename_axis("date").to_csv(OUT / "data.csv")
    metrics = {"end": END, "run_date": str(pd.Timestamp.today().date()), "finlab_version": finlab.__version__, "designs": designs, **extras(),
               "calibration": calibration(), "coverage": [coverage(pd.Timestamp(d)) for d in ("2017-09-01", "2018-04-02", "2018-06-01", "2020-01-02", "2024-01-02")],
               "yearly_funds": {t: yearly(bench(t, table.index)) for t in FUNDS[:-1]}}
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    for name, result in designs.items():
        print(name, result["original_stats"]["sharpe"], result["grid_wins"], "/", result["grid_size"], result["passes"])


if __name__ == "__main__":
    main()