"""Leveraged ETF timing: seven simple rules, three robustness checks (FinLab).

The seven rules:

  tqqq_200ma        hold TQQQ while QQQ is above its 200-day moving average, else T-bills (BIL)
  qld_200ma         the same rule with QLD (2x Nasdaq-100)
  upro_200ma        hold UPRO (3x S&P 500) while SPY is above its 200-day average, else BIL
  soxl_200ma        hold SOXL (3x semiconductors) while SOXX is above its 200-day average, else BIL
  qqq_golden_cross  hold QQQ while its 50-day average is above its 200-day average, else BIL
  spy_golden_cross  the same golden cross rule on SPY
  upro_tmf          55% UPRO / 45% TMF (3x long Treasuries), rebalanced every quarter

The 200-day rules use a 1% band: switch in when the index closes more than 1% above
its average, switch out when it closes more than 1% below, otherwise keep the current
holding. Signals are read at the close and traded at the next day's close.

"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 rule passes only if
  1. at least 7 of its 9 neighbouring parameter sets beat QQQ,
  2. it beats QQQ in both halves of the backtest, and
  3. it still beats QQQ after paying 0.1% per buy and per sell (finlab's US default
     is zero commission, so this charge stands in for spreads and slippage).

Complete the AI-assisted setup flow at https://finlab.finance/en/setup before running this file.

    python research.py                # three checks for all seven rules
    python research.py tqqq_200ma     # backtest one rule and open its report
"""
import sys

import pandas as pd
import finlab
from finlab import data
from finlab.backtest import sim

START = "2017-01-03"   # prices start in January 2016; every average up to 250 days is ready by now
CASH = "BIL"           # 1-3 month Treasury bills
COST_PER_SIDE = 0.001  # check 3: 0.1% per buy and per sell

RULES = {
    "tqqq_200ma": {"kind": "trend", "signal": "QQQ", "hold": "TQQQ",
                   "params": (200, 0.01), "grid": ((150, 200, 250), (0.0, 0.01, 0.02))},
    "qld_200ma": {"kind": "trend", "signal": "QQQ", "hold": "QLD",
                  "params": (200, 0.01), "grid": ((150, 200, 250), (0.0, 0.01, 0.02))},
    "upro_200ma": {"kind": "trend", "signal": "SPY", "hold": "UPRO",
                   "params": (200, 0.01), "grid": ((150, 200, 250), (0.0, 0.01, 0.02))},
    "soxl_200ma": {"kind": "trend", "signal": "SOXX", "hold": "SOXL",
                   "params": (200, 0.01), "grid": ((150, 200, 250), (0.0, 0.01, 0.02))},
    "qqq_golden_cross": {"kind": "cross", "signal": "QQQ", "hold": "QQQ",
                         "params": (50, 200), "grid": ((40, 50, 60), (150, 200, 250))},
    "spy_golden_cross": {"kind": "cross", "signal": "SPY", "hold": "SPY",
                         "params": (50, 200), "grid": ((40, 50, 60), (150, 200, 250))},
    "upro_tmf": {"kind": "mix", "hold": "UPRO", "hedge": "TMF",
                 "params": (0.55, "Q"), "grid": ((0.45, 0.55, 0.65), ("M", "Q", "6M"))},
}

_prices = None


def prices():
    """Dividend-adjusted closes of every ETF used here, each on its own trading days."""
    global _prices
    if _prices is None:
        data.set_market("us_fund")
        _prices = data.get("us_fund_price:adj_close")
    return _prices


def price(ticker):
    return prices()[ticker].dropna()


# ---------------------------------------------------------------------------
# Signals
# ---------------------------------------------------------------------------
def trend_state(close, window, band):
    """1 above the moving average (with a band to cut whipsaws), 0 below."""
    average = close.rolling(window).mean()
    state = pd.Series(float("nan"), index=close.index)
    state[close > average * (1 + band)] = 1.0
    state[close < average * (1 - band)] = 0.0
    return state.ffill().fillna(0.0)


def cross_state(close, short, long):
    """1 while the short moving average is above the long one (golden cross)."""
    return (close.rolling(short).mean() > close.rolling(long).mean()).astype(float)


def rebalance_dates(index, freq):
    """The first trading day, then the last trading day of every month, quarter or half-year."""
    months = {"M": 1, "Q": 3, "6M": 6}[freq]
    period = index.year * 12 + (index.month - 1) // months
    ends = pd.Series(index, index=index).groupby(period).max()
    return pd.DatetimeIndex([index[0], *ends]).unique()


# ---------------------------------------------------------------------------
# Backtest
# ---------------------------------------------------------------------------
def backtest(name, a=None, b=None, fee_ratio=None):
    rule = RULES[name]
    if a is None:
        a, b = rule["params"]

    if rule["kind"] == "mix":
        days = price(rule["hold"]).loc[START:].index
        position = pd.DataFrame({rule["hold"]: a, rule["hedge"]: 1 - a}, index=days)
        resample = rebalance_dates(days, b)
    else:
        make_state = trend_state if rule["kind"] == "trend" else cross_state
        state = make_state(price(rule["signal"]), a, b).loc[START:]
        position = pd.DataFrame({rule["hold"]: state, CASH: 1 - state})
        resample = None  # trade whenever the signal flips

    data.set_market("us_fund")
    return sim(position, resample=resample, trade_at_price="close",
               fee_ratio=fee_ratio, upload=False)


def daily_sharpe(curve):
    returns = curve.pct_change().dropna()
    return returns.mean() / returns.std() * 252 ** 0.5


def vs_qqq(curve):
    """Daily Sharpe of the curve and of QQQ (dividends included) on the same dates."""
    qqq = price("QQQ").reindex(curve.index).ffill()
    return daily_sharpe(curve), daily_sharpe(qqq)


def check(curve):
    strategy, qqq = vs_qqq(curve)
    return {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
            "sharpe": round(strategy, 4), "qqq_sharpe": round(qqq, 4), "beats_qqq": bool(strategy > qqq)}


def validate(name):
    rule = RULES[name]

    # 1. Neighbouring parameters: the original value and one step either side.
    grid = []
    for a in rule["grid"][0]:
        for b in rule["grid"][1]:
            grid.append({"params": [a, b], **check(backtest(name, a, b).creturn)})
    wins = sum(row["beats_qqq"] for row in grid)

    # 2. First half and second half of the backtest, each on its own.
    curve = backtest(name).creturn
    middle = curve.index[0] + (curve.index[-1] - curve.index[0]) / 2
    halves = [check(curve.loc[:middle]), check(curve.loc[middle:])]

    # 3. Trading costs on top of finlab's zero-commission US default.
    cost = check(backtest(name, fee_ratio=COST_PER_SIDE).creturn)

    passed = {"neighbours": wins >= 7, "halves": all(h["beats_qqq"] for h in halves),
              "costs": cost["beats_qqq"]}
    return {"rule": name, "params": list(rule["params"]), "full": check(curve),
            "grid": grid, "grid_wins": wins, "halves": halves, "cost": cost,
            "passed": {**passed, "all": all(passed.values())}}


def print_validation(result):
    full, h1, h2, cost = result["full"], *result["halves"], result["cost"]
    verdict = "PASS" if result["passed"]["all"] else "fail"
    print(f"{result['rule']:<17} {verdict:<4}  Sharpe {full['sharpe']:.2f} vs QQQ {full['qqq_sharpe']:.2f}"
          f" | neighbours {result['grid_wins']}/9"
          f" | halves {h1['sharpe']:.2f}/{h2['sharpe']:.2f} vs {h1['qqq_sharpe']:.2f}/{h2['qqq_sharpe']:.2f}"
          f" | 0.1% costs {cost['sharpe']:.2f}")


def main():
    finlab.login()  # finlab guides you through login automatically
    names = sys.argv[1:]
    if names:
        for name in names:
            report = backtest(name)
            print(name, report.get_stats())
            report.display()
        return
    for name in RULES:
        print_validation(validate(name))


if __name__ == "__main__":
    main()
