"""US earnings-surprise drift (PEAD) breadth strategy (FinLab).

Count how many liquid US stocks beat analyst EPS estimates by more than 5%
with a positive same-day price reaction over the last 20 trading days. When
that count exceeds 1% of the liquid universe AND QQQ is in an uptrend, hold
the two leveraged growth ETFs (TQQQ / TECL). Otherwise rotate into the
strongest defensive asset (IEF / GLD / SHY). When the two sleeves overlap on
a rebalance date the simulator normalizes the book (half defensive, a quarter
in each leveraged ETF). Monthly rebalance, 8% intramonth stop-loss.

pip install finlab
"""

import finlab
from finlab import data
from finlab.backtest import sim
from finlab.dataframe import FinlabDataFrame

finlab.login()  # opens a browser prompt; no token needed in the script

RISK_ASSETS = ["TQQQ", "TECL"]
DEFENSIVE_ASSETS = ["IEF", "GLD", "SHY"]

data.set_market("us")
close = data.get("us_price:adj_close")
volume = data.get("us_price:volume")
eps_actual = data.get("us_earnings_surprises:eps_actual")
eps_estimated = data.get("us_earnings_surprises:eps_estimated")
fund_close = data.get("us_fund_price:adj_close")[["QQQ"] + RISK_ASSETS + DEFENSIVE_ASSETS]

# Earnings surprise in % of the estimate (values exist on announcement days).
surprise = (eps_actual - eps_estimated) / eps_estimated.abs()

# Liquid universe: top 500 by 60-day average dollar volume, price above $5,
# excluding names with a recent extreme (>50%) daily move.
dollar_volume = (close * volume).rolling(60, min_periods=20).mean()
returns = close.pct_change()
had_extreme_move = (returns.abs() >= 0.50).rolling(252, min_periods=1).max().shift(1).fillna(False) > 0
universe = (
    dollar_volume.is_largest(500)
    & (dollar_volume > 8_000_000)
    & (close > 5)
    & ~had_extreme_move.astype(bool)
)

# Signal: breadth of positive surprise events over the last 20 trading days.
reaction = close / close.shift(1) - 1
positive_events = universe & (surprise > 0.05) & (reaction > 0)
event_count = positive_events.rolling(20, min_periods=1).sum().sum(axis=1)
event_breadth = event_count / universe.sum(axis=1).replace(0, float("nan"))

# Market regime gate: QQQ above its 200-day average with positive 6-month momentum.
qqq_trend = fund_close["QQQ"] > fund_close["QQQ"].rolling(200, min_periods=100).mean()
qqq_momentum = fund_close["QQQ"] / fund_close["QQQ"].shift(126) - 1
risk_on = (event_breadth > 0.01) & (qqq_trend & (qqq_momentum > 0)).reindex(event_breadth.index)

# Execution sleeve: two strongest leveraged ETFs when risk-on,
# strongest defensive asset when risk-off.
risk_score = FinlabDataFrame(fund_close[RISK_ASSETS].pct_change(126))
defensive_score = FinlabDataFrame(
    fund_close[DEFENSIVE_ASSETS].pct_change(63) - fund_close[DEFENSIVE_ASSETS].pct_change(21)
)
risk_on_daily = risk_on.reindex(fund_close.index).fillna(False)

risky_position = risk_score.is_largest(2) & risk_on_daily.to_frame().reindex(risk_score.index).iloc[:, 0]
risky_position = risky_position.astype(float).div(risky_position.astype(float).sum(axis=1), axis=0).fillna(0)
defensive_position = (
    defensive_score.is_largest(1)
    & (~risk_on_daily).to_frame().reindex(defensive_score.index).iloc[:, 0]
).astype(float)

position = (
    risky_position.reindex(fund_close.index).fillna(0)
    .join(defensive_position.reindex(fund_close.index).fillna(0), how="outer")
    .fillna(0)
    .T.groupby(level=0).max().T
)
position = position.loc["2016-01-01":]

data.set_market("us_fund")
report = sim(
    position,
    resample="M",
    name="us_earnings_surprise_drift",
    upload=False,
    fee_ratio=0,
    tax_ratio=0,
    trade_at_price="close",
    position_limit=1,
    stop_loss=0.08,
    touched_exit=True,
)
print(report.get_stats())
report.display()