"""
Bollinger Bands on S&P 500 stocks and on QQQ, 2016 to 2026 (FinLab research, 2026-11).

The bands are an N-day moving average (the middle band) plus and minus k population standard
deviations of the last N closes. %B is where the close sits between the bands (0 = lower band,
1 = upper band) and BandWidth is the distance between the bands divided by the middle band.

Designs, each run on a grid of N 15/20/25 days x k 1.5/2/2.5 (original setting 20 days, k = 2)
  sp500_reversion   Buy an S&P 500 member when it closes below its lower band (%B < 0). Sell when it
                    closes at or above the middle band (%B >= 0.5) or leaves the index.
  sp500_squeeze     Buy a member when it closes above its upper band and its BandWidth was at a
                    126-day low on at least one of the previous 20 days (the squeeze). Sell when
                    it closes below the middle band or leaves the index.
  qqq_reversion     sp500_reversion's rule on QQQ alone: fully in QQQ while held, cash otherwise.
  qqq_squeeze       sp500_squeeze's rule on QQQ alone.

Signals are read at a close and traded at the next day's open. Prices are FinLab's adjusted
prices. S&P 500 designs use each day's members, one share class per company, equal weight with
at most 5% in a stock, so they hold cash when fewer than 20 stocks qualify.

"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 the same exposure held fixed: QQQ (or RSP, the equal-weight
     S&P 500 fund, for the stock designs) at the design's average invested share, the rest in
     cash, rebalanced monthly,
  6. its original setting beats QQQ on returns above T-bills (BIL), because it holds cash.

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 = "2016-09-01"
END = "2026-11-05"
COST_PER_SIDE = 0.001
MAX_WEIGHT = 0.05
SQUEEZE_LOOKBACK = 126   # BandWidth at its lowest of this many days is a squeeze
SQUEEZE_RECENT = 20      # a breakout counts if a squeeze happened in this many days before it
FUNDS = ("QQQ", "SPY", "RSP")
OUT = Path(__file__).resolve().parent

DAYS = (15, 20, 25)
WIDTHS = (1.5, 2.0, 2.5)
ORIGINAL = (20, 2.0)
DESIGNS = ("sp500_reversion", "sp500_squeeze", "qqq_reversion", "qqq_squeeze")

data.truncate_end = END


@cache
def get(key, market):
    data.set_market(market)
    return data.get(key)


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


@cache
def funds():
    """Fund prices with dividends on stock-market trading days."""
    return get("us_fund_price:adj_close", "us_fund")[[*FUNDS, "BIL"]].reindex(prices().index)


@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", "us")
    profile = pd.DataFrame(get("us_company_profile", "us")).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)


def closes(design):
    return funds()[["QQQ"]] if design.startswith("qqq") else prices()[members().columns]


def bands(close, days, width):
    middle = close.rolling(days).mean()
    spread = width * close.rolling(days).std(ddof=0)
    return middle - spread, middle, middle + spread


def signals(design, days, width):
    """Entry and exit days of a design before index membership is applied."""
    close = closes(design)
    lower, middle, upper = bands(close, days, width)
    if design.endswith("reversion"):
        return close < lower, close >= middle
    bandwidth = (upper - lower) / middle
    squeeze = bandwidth <= bandwidth.rolling(SQUEEZE_LOOKBACK).min()
    recent = squeeze.astype(float).shift(1).rolling(SQUEEZE_RECENT, min_periods=1).max() > 0
    return (close > upper) & recent, close < middle


@cache
def position(design, days, width):
    entries, exits = (FinlabDataFrame(s) for s in signals(design, days, width))
    if design.startswith("sp500"):
        entries, exits = entries & members(), exits | ~members()
    return entries.hold_until(exits).loc[START:]


def run_sim(table, market, fee_ratio=None, **options):
    data.set_market(market)
    report = sim(table, trade_at_price="open", fee_ratio=fee_ratio, upload=False, **options)
    curve = report.creturn
    curve = curve[curve.index.isin(prices().index)].loc[:END]
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):], report


def run(design, setting, fee_ratio=None):
    table = position(design, *setting)
    if design.startswith("qqq"):
        return run_sim(table.astype(float), "us_fund", fee_ratio, resample=None)
    return run_sim(table.loc[:, table.any()], "us", fee_ratio, resample=None, position_limit=MAX_WEIGHT)


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


def invested(design):
    """Daily share of the account the original setting puts to work."""
    held = position(design, *ORIGINAL).sum(axis=1)
    return held.clip(upper=1) if design.startswith("qqq") else (held * MAX_WEIGHT).clip(upper=1)


def monthly(index):
    """Rebalance dates: the first day, then the last trading day of each month."""
    ends = index.to_series().groupby(index.to_period("M")).max()
    return pd.DatetimeIndex([index[0], *ends[ends > index[0]]])


@cache
def fixed_curve(design):
    """The same exposure held fixed: the fund at the design's average invested share, monthly."""
    fund = "QQQ" if design.startswith("qqq") else "RSP"
    share = invested(design).mean()
    table = pd.DataFrame({fund: share}, index=position(design, *ORIGINAL).index)
    return run_sim(table, "us_fund", resample=monthly(table.index))[0], round(float(share), 6)


# ---------------------------------------------------------------------------
# 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 trades(design):
    table = run(design, ORIGINAL)[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()), "mean_days": round(table["period"].mean(), 2)}


def turnover(design):
    """One-way turnover a year of the original setting's target weights, ignoring drift between trades."""
    table = position(design, *ORIGINAL).astype(float)
    if design.startswith("sp500"):
        table = table.mul((1 / table.sum(axis=1).clip(lower=1)).clip(upper=MAX_WEIGHT), axis=0)
    years = (table.index[-1] - table.index[0]).days / 365.25
    return round(float(table.diff().abs().sum(axis=1).sum() / 2 / years), 4)


def study(design):
    grid = {str(s): stats(backtest(design, s)) for s in product(DAYS, WIDTHS)}
    curve = backtest(design, ORIGINAL)
    wins = sum(cell["beats_qqq"] for cell in grid.values())
    first, second = halves(curve)
    fixed, share = fixed_curve(design)
    fixed = fixed.reindex(curve.index).dropna()
    held = position(design, *ORIGINAL).sum(axis=1)
    result = {
        "name": design, "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(design, 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),
        "fixed": {"invested": share, "sharpe": round(daily_sharpe(fixed.pct_change().dropna()), 6),
                  "cagr": round(cagr(fixed), 6), "mdd": round((fixed / fixed.cummax() - 1).min(), 6)},
        "invested": share,
        "days_held": round(float((held > 0).mean()), 6),
        "holdings": {"mean": round(held.mean(), 2), "median": float(held.median()), "max": int(held.max()),
                     "max_date": str(held.idxmax().date()), "days_under_20": int((held < 20).sum())},
        "turnover": turnover(design),
        "trades": trades(design),
        "rsp_wins": int(sum(c["sharpe"] > c["rsp_sharpe"] for c in grid.values())),
    }
    result["checks"] = {
        "grid": wins >= 0.7 * len(grid), "halves": first["beats_qqq"] and second["beats_qqq"],
        "cost": result["original_stats"]["beats_qqq"] and result["with_cost"]["beats_qqq"],
        "without_best_year": result["without_best_year"]["sharpe"] > result["without_best_year"]["qqq_sharpe"],
        "fixed": result["original_stats"]["sharpe"] > result["fixed"]["sharpe"],
        "above_tbills": result["excess_sharpe"] > result["qqq_excess_sharpe"],
    }
    result["checks"] = {k: bool(v) for k, v in result["checks"].items()}
    result["passes"] = all(result["checks"].values())
    return result


def main():
    results = {design: study(design) for design in DESIGNS}
    curves = {design: backtest(design, ORIGINAL) for design in DESIGNS}
    first = min(c.index[0] for c in curves.values())
    metrics = {"start": START, "end": END, "run_date": str(pd.Timestamp.now(tz="Asia/Taipei").date()),
               "finlab_version": finlab.__version__, "designs": results,
               "yearly_funds": {t: yearly(bench(t, prices().loc[first:END].index)) for t in FUNDS}}
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    table = pd.DataFrame(curves).loc[first:].ffill()
    for ticker in FUNDS:
        table[ticker] = bench(ticker, table.index)
    (table / table.bfill().iloc[0]).round(6).rename_axis("date").to_csv(OUT / "data.csv")
    for design, result in results.items():
        print(design, result["original_stats"]["sharpe"], "QQQ", result["original_stats"]["qqq_sharpe"],
              result["grid_wins"], "/", result["grid_size"], result["passes"])


if __name__ == "__main__":
    main()
