"""
Low-volatility rules on S&P 500 stocks, 2017 to 2026 (FinLab research, 2026-11).

Volatility is the standard deviation of daily returns. Designs:

  lowvol_splv          the n lowest w-day volatilities, weighted inversely to volatility,
                       after the third Friday of February, May, August and November (SPLV's rule)
                       grid: n 50/100/150 x w 126/252/378, original 100, 252
  lowvol_conservative  Blitz & van Vliet (2018): among the lowest v of 36-month volatilities,
                       the n best on the average rank of 12-1 momentum and 12-month net payout
                       yield (dividends plus buybacks minus issuance), equal weight, quarterly
                       grid: v 40/50/60% x n 40/50/60, original 50%, 50
                       36 months of prices run out in 2019-01, so it starts there
  lowvol_buyback       the n best on a blend of low w-day volatility and a four-quarter fall in
                       diluted shares (a move above 30% does not count), monthly, equal weight
                       grid: n 15/25/40 x w 126/252/378 x volatility share 0.3/0.5/0.7,
                       original 25, 252, 0.5

lowvol_buyback comes from a research notebook whose result we had already seen, so its settings
were kept and the article says so. See notebook_rule() for what was wrong with that result.

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)

WINDOWS = (126, 252, 378)

DESIGNS = {  # name: (grid, original setting)
    "lowvol_splv": (((50, 100, 150), WINDOWS), (100, 252)),
    "lowvol_conservative": (((0.4, 0.5, 0.6), (40, 50, 60)), (0.5, 50)),
    "lowvol_buyback": (((15, 25, 40), WINDOWS, (0.3, 0.5, 0.7)), (25, 252, 0.5)),
}
COMBINES = {"lowvol_conservative": "lowvol_splv", "lowvol_buyback": "lowvol_splv"}


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


# ---------------------------------------------------------------------------
# Low volatility
# ---------------------------------------------------------------------------
@cache
def daily_returns():
    return prices()[members().columns].pct_change(fill_method=None)


@cache
def volatility(window):
    return daily_returns().rolling(window, min_periods=int(window * 0.9)).std()


def third_fridays():
    days = prices().index
    fridays = pd.date_range(days[0], END, freq="WOM-3FRI")
    return pd.DatetimeIndex([days[days <= d][-1] for d in fridays if d.month in (2, 5, 8, 11)])


def splv_weights(n, window):
    table = volatility(window)[members()]
    table = table.loc[table.index.isin(third_fridays())].loc["2017-08-01":]
    picks = FinlabDataFrame(table).is_smallest(n)
    weights = (1 / table)[picks]
    return weights.div(weights.sum(axis=1), axis=0).fillna(0)


@cache
def net_payout_yield():
    paid = -(trailing("us_cash_flow:net_common_stock_issuance").fillna(0)
             + trailing("us_cash_flow:common_dividends_paid").fillna(0))
    value = market_value()
    return on_days(paid, 0) / value[value > 0]


@cache
def momentum():
    close = prices()[members().columns]
    return close.shift(21) / close.shift(252) - 1


@cache
def buyback():
    """Fall in diluted shares over four quarters. A move above 30% is a split or a merger."""
    shares = quarterly("us_income_statement:weighted_average_shs_out_dil")
    change = shares / shares.shift(4) - 1
    return on_days((-change)[change.abs() <= 0.3], -1)


@cache
def position(name, *setting):
    if name == "lowvol_conservative":
        share, n = setting
        calm = volatility(756)[members()].rank(axis=1, pct=True) <= share
        momentum_rank = momentum()[calm].rank(axis=1, pct=True)
        payout_rank = net_payout_yield()[calm].rank(axis=1, pct=True)
        return FinlabDataFrame((momentum_rank + payout_rank) / 2).is_largest(n)
    n, window, weight = setting
    calm = FinlabDataFrame((-daily_returns().rolling(window).std())[members()]).cs.rank()
    change = buyback()[buyback() > -1]
    fewer = FinlabDataFrame(change[members()]).cs.rank()
    return (calm * weight + fewer * (1 - weight))[buyback() > 0].is_largest(n)


@cache
def backtest(name, setting, fee_ratio=None):
    if name == "lowvol_splv":
        return simulate(splv_weights(*setting), third_fridays(), 1, fee_ratio)
    first = "2019-01-01" if name == "lowvol_conservative" else "2017-08-01"
    return simulate(position(name, *setting).loc[first:], "Q" if name == "lowvol_conservative" else "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"]
    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 quintiles():
    """Monthly, equal weight: each row is a fifth of the S&P 500 by trailing 252-day volatility."""
    rank = volatility(252)[members()].rank(axis=1, pct=True)
    out = {}
    for q in range(5):
        picks = ((rank > q / 5) & (rank <= (q + 1) / 5)).loc["2017-08-01":]
        curve, _ = simulate(picks, "M", 1)
        out[str(q + 1)] = {**stats(curve), "held": round(picks.sum(axis=1).resample("ME").last().mean(), 1)}
    return out


def notebook_rule():
    """The notebook's first version, as written, and with its non-trading rows removed.

    Its share table carries filing dates, so the position table gained rows that match no trading
    day. Those rows were empty, and a monthly rebalance reads the last row of the month, so whenever
    that row was one of them the backtest sat in cash for a month. That includes March 2020."""
    close = prices()
    index = get("us_index_constituents:sp500").reindex(close.index, method="ffill").eq(True)
    shares = get("us_income_statement:weighted_average_shs_out_dil")
    vol = close.pct_change(fill_method=None).rolling(252).std()
    change = shares / shares.shift(4) - 1
    bought = (-change).where(change.abs() <= 0.3)
    score = (-vol).where(index).cs.rank() * 0.5 + bought.where(index).cs.rank() * 0.5
    table = score.where(index & (bought > 0)).is_largest(25).loc[:END]
    as_written, report = simulate(table, "M", 1)
    fixed, _ = simulate(table.reindex(close.index), "M", 1)
    held = (report.position.fillna(0).astype(float) != 0).sum(axis=1).loc[as_written.index[1]:]
    return {"as_written": stats(as_written), "fixed": stats(fixed), "rebalances": int(len(held)),
            "empty": int((held == 0).sum()), "empty_months": [str(d.date()) for d in held[held == 0].index],
            "share_of_days_empty": round(float(as_written.pct_change().eq(0).mean()), 4),
            "curves": {"as_written": as_written, "fixed": fixed}}


def crash_and_bear(curves):
    """Return from the 2020 market top to its low, and over 2022."""
    def span(curve, a, b):
        part = curve.loc[a:b]
        return round(part.iloc[-1] / part.iloc[0] - 1, 6)
    return {k: {"covid": span(c, "2020-02-19", "2020-03-23"), "y2022": span(c, "2021-12-31", "2022-12-30")} for k, c in curves.items()}


def rule_vs_fund():
    curve = backtest("lowvol_splv", (100, 252))[0]
    table = pd.concat([curve.pct_change(), bench("SPLV", curve.index).pct_change(), bench("USMV", curve.index).pct_change()], axis=1).dropna()
    table.columns = ["rule", "SPLV", "USMV"]
    return {"corr_splv": round(table.corr().loc["rule", "SPLV"], 4), "corr_usmv": round(table.corr().loc["rule", "USMV"], 4)}


def extras():
    notebook = notebook_rule()
    curves = {name: backtest(name, DESIGNS[name][1])[0] for name in ("lowvol_splv", "lowvol_buyback")}
    curves |= {t: bench(t, curves["lowvol_splv"].index) for t in ("QQQ", "SPY", "RSP", "SPLV", "USMV")}
    notebook_curves = notebook.pop("curves")
    pd.DataFrame(notebook_curves).round(6).to_csv(OUT / "notebook_rule.csv")
    return {"quintiles": quintiles(), "notebook_rule": notebook, "stress": crash_and_bear(curves), "rule_vs_fund": rule_vs_fund()}


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