"""
Golden crosses and death crosses on S&P 500 stocks, 2017 to 2026 (FinLab research, 2026-11).

Designs, each with a grid around the original setting
  golden_cross    Buy a stock when its S-day average crosses above its L-day average and sell when
                  it crosses back or the stock leaves the index.
                  grid: S 40/50/60 x L 150/200/250
  cross_revenue   golden_cross, entries also need quarterly revenue up at least g on a year earlier.
                  grid: golden_cross x g 5/10/20%

Event study: every 50/200 golden cross and death cross of an index member from START. Returns run
from the next close for 21, 63, 126 and 252 trading days and are compared with the average
member over the same days. Events without that many days left, including stocks acquired or
delisted before then, are dropped.

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.

"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. cross_revenue also beats golden_cross at its original setting.

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
HORIZONS = (21, 63, 126, 252)   # trading days after an event
WHIPSAW_DAYS = 63
FUNDS = ("QQQ", "SPY", "RSP")
OUT = Path(__file__).resolve().parent

FAST = (40, 50, 60)
SLOW = (150, 200, 250)
REVENUE = (0.05, 0.10, 0.20)

DESIGNS = {  # name: (grid, original setting)
    "golden_cross": ((FAST, SLOW), (50, 200)),
    "cross_revenue": ((FAST, SLOW, REVENUE), (50, 200, 0.10)),
}
HOLDS_CASH = tuple(DESIGNS)
COMBINES = {"cross_revenue": "golden_cross"}

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 crosses(fast, slow):
    """Golden and death crosses: the day the fast average moves above or below the slow one."""
    a, b = average(fast), average(slow)
    return (a > b) & (a.shift(1) <= b.shift(1)), (a < b) & (a.shift(1) >= b.shift(1))


@cache
def position(name, fast, slow, *rest):
    entries = crosses(fast, slow)[0] & members()
    if rest:
        entries = entries & (growth("us_income_statement:revenue") >= rest[0])
    exits = average(fast) < average(slow)
    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


# ---------------------------------------------------------------------------
# Event study
# ---------------------------------------------------------------------------
def forward(days):
    price = member_prices()
    return price.shift(-days - 1) / price.shift(-1) - 1


def event_returns(events, days):
    """One row per event with its return and its return above the average member."""
    ahead = forward(days)
    market = ahead.where(members()).mean(axis=1)
    picked = ahead.where(events).loc[START:].stack().dropna()
    return pd.DataFrame({"raw": picked, "excess": picked - market.reindex(picked.index.get_level_values(0)).to_numpy()})


def beat_average_share(days):
    """How often any member beat the average member over `days`, as a base rate."""
    ahead = forward(days).where(members()).loc[START:]
    return round(float((ahead.gt(ahead.mean(axis=1), axis=0)).sum().sum() / ahead.notna().sum().sum()), 6)


def market_backdrop(events):
    """Where RSP stood against its high on event days, against all days."""
    fund = funds()["RSP"]
    below = (fund / fund.cummax() - 1).loc[START:]
    count = events.loc[START:].sum(axis=1)
    return {"down_10_pct": round(float(count[below <= -0.10].sum() / count.sum()), 6),
            "within_5_pct": round(float(count[below > -0.05].sum() / count.sum()), 6)}


def summary(rows):
    return {"events": len(rows), "mean": round(rows["raw"].mean(), 6), "median": round(rows["raw"].median(), 6),
            "positive": round((rows["raw"] > 0).mean(), 6), "excess_mean": round(rows["excess"].mean(), 6),
            "market_mean": round((rows["raw"] - rows["excess"]).mean(), 6),
            "excess_median": round(rows["excess"].median(), 6), "beat_average": round((rows["excess"] > 0).mean(), 6),
            "p10": round(rows["excess"].quantile(0.1), 6), "p90": round(rows["excess"].quantile(0.9), 6)}


def event_study():
    golden, death = crosses(50, 200)
    golden, death = golden & members(), death & members()
    table = {}
    for label, events in (("golden", golden), ("death", death)):
        rows = {days: event_returns(events, days) for days in HORIZONS}
        counts = events.loc[START:].sum(axis=1)
        table[label] = {"count": int(counts.sum()), "backdrop": market_backdrop(events),
                        "by_year": {str(y): int(n) for y, n in counts.groupby(counts.index.year).sum().items()},
                        "horizons": {str(days): summary(frame) for days, frame in rows.items()},
                        "excess_126_by_year": {str(y): round(v, 6) for y, v in rows[126]["excess"].groupby(rows[126].index.get_level_values(0).year).mean().items()}}
    a, b = average(50), average(200)
    undone = (a < b).astype(float)[::-1].rolling(WHIPSAW_DAYS, min_periods=WHIPSAW_DAYS).max()[::-1].shift(-1)
    table["golden_undone_within_63_days"] = round(undone.where(golden).loc[START:].stack().dropna().mean(), 6)
    table["any_member_beats_average"] = {str(days): beat_average_share(days) for days in HORIZONS}
    below = (funds()["RSP"] / funds()["RSP"].cummax() - 1).loc[START:]
    table["all_days_backdrop"] = {"down_10_pct": round(float((below <= -0.10).mean()), 6),
                                  "within_5_pct": round(float((below > -0.05).mean()), 6)}
    return table


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 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()))
    curves = {name: backtest(name, DESIGNS[name][1]) for name in DESIGNS}
    daily = pd.DataFrame({"golden_cross": curves["golden_cross"], "RSP": bench("RSP", curves["golden_cross"].index)}).pct_change()
    results["golden_cross"]["rsp_correlation"] = round(float(daily.corr().iloc[0, 1]), 6)
    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(), "events": event_study(),
               "yearly_funds": {t: yearly(bench(t, curves["golden_cross"].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()