"""QQQ breakout / trend-following -> leveraged-ETF strategy (FinLab, US market).

The signal lives entirely on QQQ: the strategy goes risk-on only when QQQ has
recently closed above its prior 21-day high, sits above its 200-day moving
average, and has a positive 126-day return. In risk-on it holds the two
strongest leveraged growth ETFs (TQQQ, TECL); otherwise it rotates into one
defensive ETF (IEF, GLD, or SHY). Reproduces the backtest in the article.
Complete the AI-assisted setup flow at https://finlab.finance/en/setup before running this file.

    python strategy.py
"""
import finlab
from finlab import data
from finlab.backtest import sim
from finlab.dataframe import FinlabDataFrame

RISK_ASSETS = ["TQQQ", "TECL"]            # leveraged growth ETFs (held in risk-on)
DEFENSIVE_ASSETS = ["IEF", "GLD", "SHY"]  # treasuries / gold / cash-like (risk-off)


def build_position():
    data.set_market("us_fund")
    close = data.get("us_fund_price:adj_close")[["QQQ"] + RISK_ASSETS + DEFENSIVE_ASSETS]

    # 1. Trend context: QQQ above its 200-day average with positive 126-day momentum.
    qqq = close["QQQ"]
    qqq_trend = qqq > qqq.rolling(200, min_periods=100).mean()
    qqq_momentum = qqq / qqq.shift(126) - 1

    # 2. Breakout trigger: QQQ closes above its prior 21-day high; the trigger
    #    stays "active" for 5 sessions so a monthly rebalance can catch it.
    qqq_breakout = qqq > qqq.shift(1).rolling(21, min_periods=10).max()
    breakout_active = qqq_breakout.rolling(5, min_periods=1).max().astype(bool)

    # 3. Risk-on only when breakout AND trend AND momentum all agree.
    risk_on = qqq_trend & (qqq_momentum > 0) & breakout_active

    # 4. In risk-on: hold the 2 strongest leveraged ETFs by 126-day return.
    #    In risk-off: hold the 1 defensive ETF with the best recent momentum.
    risk_score = FinlabDataFrame(close[RISK_ASSETS].pct_change(126))
    defensive_score = FinlabDataFrame(
        close[DEFENSIVE_ASSETS].pct_change(63) - close[DEFENSIVE_ASSETS].pct_change(21)
    )
    risky = risk_score.is_largest(2) & risk_on.to_frame().reindex(risk_score.index).iloc[:, 0]
    risky = risky.astype(float).div(risky.astype(float).sum(axis=1), axis=0).fillna(0)
    defensive = (
        defensive_score.is_largest(1) & (~risk_on).to_frame().reindex(defensive_score.index).iloc[:, 0]
    ).astype(float)

    position = (
        risky.reindex(close.index).fillna(0)
        .join(defensive.reindex(close.index).fillna(0), how="outer")
        .fillna(0)
        .T.groupby(level=0).max().T
    )
    return position.loc["2016-01-01":]


def main():
    finlab.login()  # finlab guides you through login automatically
    position = build_position()
    report = sim(
        position,
        resample="M",           # rebalance monthly
        position_limit=1,       # max weight PER asset (1.0 = 100%); the number of
                                # holdings is set by is_largest(2) above -> 2 ETFs at 50% each
        stop_loss=0.08,         # 8% touched stop on each leg
        touched_exit=True,
        trade_at_price="close",
        upload=False,
    )
    print(report.get_stats())
    report.display()


if __name__ == "__main__":
    main()