"""
52-week highs on S&P 500 stocks, 2017 to 2026 (FinLab research, 2026-11).

Designs, each with a grid around the original setting
  new_high          Buy a stock that closes at its highest in L days and sell when it closes
                    below its M-day average or leaves the index.
                    grid: L 200/252/300 x M 20/50/100
  new_high_revenue  new_high, entries also need quarterly revenue up at least g on a year earlier.
                    grid: new_high x g 5/10/20%
  new_high_eps      new_high, entries also need quarterly EPS up at least g on a year earlier,
                    positive in both quarters.
                    grid: new_high x g 10/20/30%

Prices are FinLab's adjusted closes. 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, equal weight with at most 5% in a stock, so the portfolio holds cash
when fewer than 20 stocks qualify. Every design starts on 2017-04-03, the first month start
with 300 days of prices.

"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. its original setting beats QQQ on returns above T-bills (BIL), because it can hold cash,
  6. new_high_revenue and new_high_eps also beat new_high at their original settings.

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

START = "2017-04-03"
END = "2026-10-30"
COST_PER_SIDE = 0.001
MAX_WEIGHT = 0.05
FUNDS = ("QQQ", "SPY", "RSP")
OUT = Path(__file__).resolve().parent

LOOKBACKS = (200, 252, 300)
AVERAGES = (20, 50, 100)
REVENUE = (0.05, 0.10, 0.20)
EPS = (0.10, 0.20, 0.30)

DESIGNS = {  # name: (grid, original setting)
    "new_high": ((LOOKBACKS, AVERAGES), (252, 50)),
    "new_high_revenue": ((LOOKBACKS, AVERAGES, REVENUE), (252, 50, 0.10)),
    "new_high_eps": ((LOOKBACKS, AVERAGES, EPS), (252, 50, 0.20)),
}
HOLDS_CASH = tuple(DESIGNS)
COMBINES = {"new_high_revenue": "new_high", "new_high_eps": "new_high"}
GROWTH = {"new_high_revenue": "us_income_statement:revenue", "new_high_eps": "us_income_statement:eps_diluted"}

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")[[*FUNDS, "BIL"]].reindex(prices().index)


# ---------------------------------------------------------------------------
# S&P 500 signals
# ---------------------------------------------------------------------------
@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")
    profile = pd.DataFrame(get("us_company_profile")).astype({"stock_id": str}).drop_duplicates("stock_id")
    names = profile.set_index("stock_id")["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 member_prices():
    return prices()[members().columns]


@cache
def average(days):
    return member_prices().rolling(days).mean()


@cache
def growth(key):
    """Year-over-year growth of a quarterly item on the trading-day grid, from its filing date.
    Quarters without a growth rate become -100%, or they would inherit the previous quarter's."""
    quarters = get("us_balance_sheet:total_assets").index
    item = get(key).reindex(index=quarters, columns=members().columns)
    year_ago = item.shift(4)
    change = (item / year_ago[year_ago > 0] - 1)[item > 0].fillna(-1)
    return change.index_str_to_date().reindex(index=prices().index, method="ffill")


def new_highs(lookback):
    price = member_prices()
    return (price >= price.rolling(lookback).max()) & members()


@cache
def position(name, lookback, days, *rest):
    entries = new_highs(lookback)
    if rest:
        entries = entries & (growth(GROWTH[name]) >= rest[0])
    exits = member_prices() < average(days)
    return entries.hold_until(exits | ~members()).loc[START:]


def run(name, setting, fee_ratio=None):
    table = position(name, *setting)
    report = sim(table.loc[:, table.any()], resample=None, trade_at_price="close", fee_ratio=fee_ratio,
                 position_limit=MAX_WEIGHT, market="us", upload=False)
    curve = report.creturn
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):], report


@cache
def backtest(name, setting, fee_ratio=None):
    return run(name, setting, fee_ratio)[0]


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------
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())}
    for label, series in (("", curve), *((f"{t.lower()}_", bench(t, curve.index)) for t in FUNDS)):
        returns = series.pct_change()
        out[f"{label}sharpe"] = round(daily_sharpe(returns.dropna()), 6)
        out[f"{label}cagr"] = round(cagr(series), 6)
        out[f"{label}mdd"] = round((series / series.cummax() - 1).min(), 6)
        out[f"{label}volatility"] = round(returns.std() * 252 ** 0.5, 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)) for s in product(*axes)}
    curve = backtest(name, original)
    wins = sum(cell["beats_qqq"] for cell in grid.values())
    first, second = halves(curve)
    result = {
        "name": name, "original": str(original), "original_stats": grid[str(original)], "grid": grid,
        "grid_wins": wins, "grid_size": len(grid), "first_half": first, "second_half": second,
        "with_cost": stats(backtest(name, original, COST_PER_SIDE)), "yearly": yearly(curve),
        "without_best_year": without_best_year(curve), "excess_sharpe": round(excess_sharpe(curve), 6),
        "qqq_excess_sharpe": round(excess_sharpe(bench("QQQ", curve.index)), 6),
    }
    result["passes"] = bool(
        wins >= 0.7 * len(grid) and first["beats_qqq"] and second["beats_qqq"]
        and result["original_stats"]["beats_qqq"] and result["with_cost"]["beats_qqq"]
        and result["without_best_year"]["sharpe"] > result["without_best_year"]["qqq_sharpe"]
        and (name not in HOLDS_CASH or result["excess_sharpe"] > result["qqq_excess_sharpe"])
        and (name not in COMBINES or grid[str(original)]["sharpe"] > stats(backtest(COMBINES[name], DESIGNS[COMBINES[name]][1]))["sharpe"]))
    return result


def trades(name):
    table = run(name, DESIGNS[name][1])[1].get_trades().dropna(subset=["exit_date"])
    return {"count": len(table), "win_rate": round((table["return"] > 0).mean(), 6),
            "mean_return": round(table["return"].mean(), 6), "median_return": round(table["return"].median(), 6),
            "median_days": float(table["period"].median())}


def turnover(table):
    """One-way turnover a year of the target weights, ignoring drift between trades."""
    weights = table.astype(float).mul((1 / table.sum(axis=1).clip(lower=1)).clip(upper=MAX_WEIGHT), axis=0)
    return round(float(weights.diff().abs().sum(axis=1).sum() / 2 / ((table.index[-1] - table.index[0]).days / 365.25)), 4)


def first_day_members():
    listed = get("us_index_constituents:sp500").loc[START]
    listed = listed[listed.astype(bool)].index
    return {"members": len(listed), "priced": int(prices().reindex(columns=listed).loc[START].notna().sum())}


def above_base(name, results):
    """Settings of a growth design above the plain rule with the same lookback and exit."""
    base = {eval(k): v["sharpe"] for k, v in results["new_high"]["grid"].items()}
    return sum(cell["sharpe"] > base[eval(k)[:2]] for k, cell in results[name]["grid"].items())


def worst_stretch(curve):
    """The maximum drawdown's peak and trough, and what RSP and QQQ did between them."""
    trough = (curve / curve.cummax() - 1).idxmin()
    peak = curve.loc[:trough].idxmax()
    out = {"peak": str(peak.date()), "trough": str(trough.date())}
    for ticker in ("RSP", "QQQ"):
        out[ticker] = round(float(funds()[ticker].loc[trough] / funds()[ticker].loc[peak] - 1), 6)
    return out


def after_low(curve, start="2020-03-23", end="2020-06-30"):
    """The rebound from the March 2020 low, when few stocks were near a high."""
    held = position("new_high", *DESIGNS["new_high"][1]).sum(axis=1)
    return {"start": start, "end": end, "rule": round(float(curve.loc[end] / curve.loc[start] - 1), 6),
            "rsp": round(float(funds()["RSP"].loc[end] / funds()["RSP"].loc[start] - 1), 6),
            "mean_holdings": round(float(held.loc[start:end].mean()), 2)}


def main():
    results = {name: study(name) for name in DESIGNS}
    for name, result in results.items():
        held = position(name, *DESIGNS[name][1]).sum(axis=1)
        result["holdings"] = {"mean": round(held.mean(), 2), "median": float(held.median()),
                              "min": int(held.min()), "max": int(held.max())}
        result["invested"] = round((held * MAX_WEIGHT).clip(upper=1).mean(), 6)
        result["turnover"] = turnover(position(name, *DESIGNS[name][1]))
        result["days_under_20_in_2017"] = int((held.loc[:"2017-12-31"] < 20).sum())
        result["first_day_holdings"] = int(held.iloc[0])
        result["trades"] = trades(name)
        for fund in ("rsp", "spy"):
            result[f"{fund}_wins"] = int(sum(c["sharpe"] > c[f"{fund}_sharpe"] for c in result["grid"].values()))
    highs = new_highs(252).loc[START:].sum(axis=1)
    curves = {name: backtest(name, DESIGNS[name][1]) for name in DESIGNS}
    for name in COMBINES:
        results[name]["above_base"] = int(above_base(name, results))
    results["new_high"]["worst_stretch"] = worst_stretch(curves["new_high"])
    results["new_high"]["after_2020_low"] = after_low(curves["new_high"])
    metrics = {
        "end": END, "run_date": str(pd.Timestamp.now(tz="Asia/Taipei").date()), "finlab_version": finlab.__version__,
        "designs": results, "first_day_members": first_day_members(),
        "new_high_days": {"mean": round(highs.mean(), 2), "median": float(highs.median()), "max": int(highs.max()),
                          "max_day": str(highs.idxmax().date()), "days_under_5": round(float((highs < 5).mean()), 6),
                          "monthly_mean": {str(d.date()): round(v, 2) for d, v in highs.resample("ME").mean().items()}},
        "yearly_funds": {t: yearly(bench(t, curves["new_high"].index)) for t in FUNDS},
    }
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    table = pd.DataFrame(curves)
    for ticker in FUNDS:
        table[ticker] = bench(ticker, table.index)
    (table / table.iloc[0]).round(6).rename_axis("date").to_csv(OUT / "data.csv")
    for name, result in results.items():
        print(name, result["original_stats"]["sharpe"], result["grid_wins"], "/", result["grid_size"], result["passes"])


if __name__ == "__main__":
    main()