"""
PEG ratio rules on S&P 500 stocks, 2018 to 2026 (FinLab research, 2026-11).

PEG is the P/E divided by EPS growth in percent. P/E is market value over four quarters of net
income. EPS is four quarters of net income per diluted share, positive at both ends. Designs:

  peg_under_one   every stock with PEG below t and one-year EPS growth above 0 and at most cap
                  grid: t 0.75/1.0/1.25 x cap 50/100/200%, original 1.0, 100%
  peg_multi_year  the same with y-year annualised EPS growth, at most 100%
                  grid: t 0.75/1.0/1.25 x y 2/3/4, original 1.0, 3
                  y needs y more years of filings, so y = 2, 3, 4 start in 2019-03, 2020-03, 2021-03,
                  the first March about 70% of members have a growth rate
  peg_in_sector   the n lowest PEGs in each of the 11 sectors, one-year growth at most cap
                  grid: n 2/3/4 x cap 50/100/200%, original 3, 100%

Each holds equal weights with at most 5% in a stock and changes hands monthly. The others start
in 2018-03, the first month 70% of members have a growth rate.

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)

THRESHOLDS = (0.75, 1.0, 1.25)
CAPS = (0.5, 1.0, 2.0)

DESIGNS = {  # name: (grid, original setting)
    "peg_under_one": ((THRESHOLDS, CAPS), (1.0, 1.0)),
    "peg_multi_year": ((THRESHOLDS, (2, 3, 4)), (1.0, 3)),
    "peg_in_sector": (((2, 3, 4), CAPS), (3, 1.0)),
}


@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)


# ---------------------------------------------------------------------------
# PEG
# ---------------------------------------------------------------------------
@cache
def trailing_eps():
    """Four quarters of net income per diluted share. eps_diluted itself is rounded after splits
    are restated: NVDA reads 0.023 for a quarter in 2019."""
    income = quarterly("us_income_statement:net_income")
    shares = quarterly("us_income_statement:weighted_average_shs_out_dil")
    return (income[income != 0] / shares[shares > 0]).rolling(4, min_periods=4).sum()


@cache
def eps_growth(years):
    eps = trailing_eps()
    earlier = eps.shift(4 * years)
    return (eps[eps > 0] / earlier[earlier > 0]) ** (1 / years) - 1


@cache
def price_to_earnings():
    income = on_days(trailing("us_income_statement:net_income"), -1)
    return market_value() / income[income > 0]


@cache
def growth_on_days(years):
    return on_days(eps_growth(years), -1)


@cache
def peg(years, cap):
    growth = growth_on_days(years)
    return (price_to_earnings() / (growth[(growth > 0) & (growth <= cap)] * 100))[members()]


def lowest_in_each_sector(table, n):
    groups = sectors().reindex(table.columns)
    picks = [table[groups.index[groups == s]].is_smallest(n) for s in groups.dropna().unique()]
    return pd.concat(picks, axis=1).reindex(columns=table.columns).fillna(False).astype(bool)


@cache
def position(name, *setting):
    if name == "peg_under_one":
        threshold, cap = setting
        return peg(1, cap) < threshold
    if name == "peg_multi_year":
        threshold, years = setting
        return peg(years, 1.0) < threshold
    n, cap = setting
    return lowest_in_each_sector(peg(1, cap), n)


def first_month(name, setting):
    return {2: "2019-03-01", 3: "2020-03-01", 4: "2021-03-01"}[setting[1]] if name == "peg_multi_year" else "2018-03-01"


@cache
def backtest(name, setting, fee_ratio=None):
    table = position(name, *setting).loc[first_month(name, setting):]
    return simulate(table, "M", 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"]
    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"]))
    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
# ---------------------------------------------------------------------------
BANDS = ((0, 0.5), (0.5, 1), (1, 1.5), (1.5, 2), (2, 3), (3, float("inf")))


def bands():
    """Everything monthly and equal weight, so each row is a slice of the S&P 500 by PEG."""
    table = peg(1, 1.0)
    rows = {f"{low}-{high}": (table >= low) & (table < high) for low, high in BANDS}
    rows["no_peg"] = table.isna() & members()
    out = {}
    for label, picks in rows.items():
        picks = picks.loc["2018-03-01":]
        curve, _ = simulate(picks, "M", 1)
        first, second = halves(curve)
        out[label] = {**stats(curve), "held": round(picks.sum(axis=1).resample("ME").last().mean(), 1),
                      "first_half_sharpe": first["sharpe"], "second_half_sharpe": second["sharpe"]}
    return out


def share_by_month():
    table, inside = peg(1, 1.0).loc["2018-03-01":], members().loc["2018-03-01":]
    days = table.index.to_series().groupby(table.index.to_period("M")).max()
    table, inside = table.loc[days], inside.loc[days]
    return {"share_with_peg": {str(d.date()): round(v, 4) for d, v in (table.notna().sum(axis=1) / inside.sum(axis=1)).items()},
            "share_under_one": {str(d.date()): round(v, 4) for d, v in ((table < 1).sum(axis=1) / inside.sum(axis=1)).items()},
            "median": {str(d.date()): round(v, 4) for d, v in table.median(axis=1).items()}}


def same_window():
    """Multi-year PEG: every setting from the first month the four-year version exists."""
    out = {}
    for setting in product(*DESIGNS["peg_multi_year"][0]):
        curve = backtest("peg_multi_year", setting)[0].loc["2021-04-01":]
        out[str(setting)] = {k: v for k, v in stats(curve).items() if k in ("start", "sharpe", "qqq_sharpe", "spy_sharpe", "rsp_sharpe", "cagr", "qqq_cagr")}
    return out


def by_start_month():
    """The three-year original from each monthly start, and one-year PEG below 1 from April 2020."""
    curve = backtest("peg_multi_year", (1.0, 3))[0]
    rows = {}
    for month in pd.date_range("2020-04-01", "2023-04-01", freq="MS"):
        cell = stats(curve.loc[month:])
        rows[str(month.date())[:7]] = {"sharpe": cell["sharpe"], "qqq_sharpe": cell["qqq_sharpe"]}
    cell = stats(backtest("peg_under_one", (1.0, 1.0))[0].loc["2020-04-01":])
    return {"three_year": rows, "one_year_from_2020_04": {"sharpe": cell["sharpe"], "qqq_sharpe": cell["qqq_sharpe"]}}


def extras():
    return {"bands": bands(), "by_month": share_by_month(), "multi_year_same_window": same_window(),
            "by_start_month": by_start_month()}


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()