"""QQQ volatility targeting without leverage (FinLab).

At every month end, measure QQQ's realised volatility over the last 21 trading days
(annualised). Hold QQQ at weight min(1, 15% / volatility) and put the rest in
short-term Treasuries (SHY). Calm markets keep the account fully in QQQ; when
volatility jumps, part of it moves to SHY until the next month end. No leverage:
the QQQ weight never goes above 100%.

The idea follows Moreira & Muir (2017), "Volatility-Managed Portfolios".

Complete the AI-assisted setup flow at https://finlab.finance/en/setup before running this file.
Needs finlab 2.2.21 or later (older versions could trade US funds on market holidays).

    python strategy.py              # backtest and print the statistics
    python strategy.py --validate   # the three robustness checks from the article
    python strategy.py --export     # write data.csv and metrics.json next to this file
"""
import json
import sys
from pathlib import Path

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

TARGET_VOLATILITY = 0.15   # annualised
VOLATILITY_DAYS = 21       # trading days of daily returns behind each estimate
RISKY, SAFE = "QQQ", "SHY"
START = "2017-07-31"       # first month end of the article's backtest


def prices():
    data.set_market("us_fund")
    # us_fund_price keeps market-holiday rows as blanks; drop them before rolling
    return data.get("us_fund_price:adj_close")[[RISKY, SAFE]].dropna()


def build_position(target=TARGET_VOLATILITY, days=VOLATILITY_DAYS):
    close = prices()
    volatility = close[RISKY].pct_change(fill_method=None).rolling(days).std() * 252 ** 0.5
    qqq_weight = (target / volatility).clip(upper=1)
    position = pd.DataFrame({RISKY: qqq_weight, SAFE: 1 - qqq_weight})
    return position.loc[START:]


def run_backtest(position, fee_ratio=None):
    data.set_market("us_fund")
    return sim(
        position,
        resample="M",            # weights are read on the last trading day of each month
        trade_at_price="close",
        fee_ratio=fee_ratio,     # None = finlab's US default (no commission, no tax)
        upload=False,
    )


# ---------------------------------------------------------------------------
# Robustness checks behind the article:  python strategy.py --validate
# "Beats QQQ" means a higher daily Sharpe ratio (no risk-free rate deducted) than
# buying and holding QQQ, dividends included, over exactly the same dates.
# ---------------------------------------------------------------------------
COST_PER_SIDE = 0.001  # 0.1% per buy and per sell, for spread and slippage


def equity_curve(report):
    """From the day the first trade fills."""
    curve = report.creturn
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):]


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


def summary(curve):
    qqq = prices()[RISKY].reindex(curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    out = {}
    for name, c in (("strategy", curve), ("qqq", qqq)):
        out[name] = {
            "cagr": round((c.iloc[-1] / c.iloc[0]) ** (1 / years) - 1, 6),
            "sharpe": round(daily_sharpe(c), 6),
            "mdd": round((c / c.cummax() - 1).min(), 6),
        }
    return {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date()), **out}


def average_qqq_weight(position):
    """Mean QQQ weight over the month ends that traded (the last one fills after the data ends)."""
    month_ends = position.groupby(position.index.to_period("M")).tail(1).iloc[:-1]
    return float(month_ends[RISKY].mean())


def validate():
    # 1. Neighbouring parameters: the original value and one step either side.
    wins = 0
    for target in (0.10, 0.15, 0.20):
        for days in (10, 21, 42):
            s = summary(equity_curve(run_backtest(build_position(target, days))))
            wins += s["strategy"]["sharpe"] > s["qqq"]["sharpe"]
            print(f"target {target:.0%}, {days}-day volatility: "
                  f"Sharpe {s['strategy']['sharpe']:.2f} vs QQQ {s['qqq']['sharpe']:.2f}")
    print(f"{wins}/9 parameter sets beat QQQ")

    # 2. First half and second half of the backtest, each on its own.
    curve = equity_curve(run_backtest(build_position()))
    middle = curve.index[len(curve) // 2]
    for label, part in (("first half", curve.loc[:middle]), ("second half", curve.loc[middle:])):
        s = summary(part)
        print(f"{label} {s['start']} to {s['end']}: "
              f"Sharpe {s['strategy']['sharpe']:.2f} vs QQQ {s['qqq']['sharpe']:.2f}")

    # 3. Trading costs: finlab's US default is zero, so charge 0.1% per side.
    s = summary(equity_curve(run_backtest(build_position(), fee_ratio=COST_PER_SIDE)))
    print(f"with 0.1% per side: CAGR {s['strategy']['cagr']:.1%}, "
          f"Sharpe {s['strategy']['sharpe']:.2f} vs QQQ {s['qqq']['sharpe']:.2f}")

    # Control: the same average QQQ weight held fixed. Volatility timing has to beat it.
    position = build_position()
    average = average_qqq_weight(position)
    fixed = pd.DataFrame({RISKY: average, SAFE: 1 - average}, index=position.index)
    s = summary(equity_curve(run_backtest(fixed)))
    print(f"fixed {average:.0%} QQQ / {1 - average:.0%} SHY: Sharpe {s['strategy']['sharpe']:.2f}")


def export(report):
    curve = equity_curve(report)
    qqq = prices()[RISKY].reindex(curve.index)
    frame = pd.DataFrame({"Volatility targeting (QQQ/SHY)": curve, "QQQ (buy and hold)": qqq})
    here = Path(__file__).resolve().parent
    (frame / frame.iloc[0]).round(6).to_csv(here / "data.csv", index_label="date")
    metrics = {"target_volatility": TARGET_VOLATILITY, "volatility_days": VOLATILITY_DAYS,
               **summary(curve), "average_qqq_weight": round(average_qqq_weight(build_position()), 6),
               "finlab_version": finlab.__version__}
    (here / "metrics.json").write_text(json.dumps(metrics, indent=1))


def main():
    finlab.login()  # finlab guides you through login automatically
    if "--validate" in sys.argv:
        validate()
        return
    report = run_backtest(build_position())
    if "--export" in sys.argv:
        export(report)
    print(report.get_stats())
    report.display()


if __name__ == "__main__":
    main()
