# Sector Momentum Breakout & Rotation -- Python / VectorBT
# Strategy code: data, indicators, signals, backtest execution, basic stats.
# Companion file: sector-momentum-python-analysis.py
# Full notebook: https://quantstr.at/wp-content/uploads/reports/sector-momentum-python.html

import vectorbt as vbt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

print(f"vectorbt {vbt.__version__} | pandas {pd.__version__} | numpy {np.__version__}")

# EMA and RSI are implemented directly here (rather than using a library
# default) to guarantee the exact same indicator math as the R notebook's
# TTR::EMA()/TTR::RSI(): both are SMA-seeded recursive smoothers -- EMA uses
# alpha = 2/(n+1); RSI uses Wilder's smoothing (alpha = 1/n) on average
# gain/loss.
def ema(series: pd.Series, n: int) -> pd.Series:
    x = series.to_numpy(dtype=float)
    alpha = 2.0 / (n + 1)
    out = np.full(len(x), np.nan)
    if len(x) >= n:
        out[n - 1] = np.mean(x[:n])
        for i in range(n, len(x)):
            out[i] = alpha * x[i] + (1 - alpha) * out[i - 1]
    return pd.Series(out, index=series.index)

def rsi(series: pd.Series, n: int) -> pd.Series:
    x = series.to_numpy(dtype=float)
    delta = np.diff(x)
    gain, loss = np.maximum(delta, 0), np.maximum(-delta, 0)
    alpha = 1.0 / n

    def wilder_smooth(v):
        out = np.full(len(v), np.nan)
        if len(v) >= n:
            out[n - 1] = np.mean(v[:n])
            for i in range(n, len(v)):
                out[i] = alpha * v[i] + (1 - alpha) * out[i - 1]
        return out

    avg_gain, avg_loss = wilder_smooth(gain), wilder_smooth(loss)
    rs = avg_gain / avg_loss
    out = 100 - 100 / (1 + rs)
    return pd.Series(np.concatenate([[np.nan], out]), index=series.index)

symbols = ["XLK", "XLF", "XLI", "XLE", "XLV", "XLY"]
start_date = "2023-08-29"
end_date = "2026-08-27"
init_cash = 100000

data = vbt.YFData.download(symbols, start=start_date, end=end_date)
close = data.get('Close')
open_price = data.get('Open')

# Some data providers' UTC timestamps include one extra session before the
# requested start date. Filtering to the calendar date keeps the same
# trading-day range the R notebook uses.
start_calendar_date = pd.Timestamp(start_date).date()
in_range = close.index.date >= start_calendar_date
close = close[in_range]
open_price = open_price[in_range]
print("Bars downloaded per symbol:")
print(close.count())

ema_fast = close.apply(lambda col: ema(col, 10))
ema_slow = close.apply(lambda col: ema(col, 30))
rsi14 = close.apply(lambda col: rsi(col, 14))

# Entry fires only on the bar where the compound condition (10-day EMA above
# 30-day EMA, AND RSI above 50) first becomes true -- not on every bar the
# condition merely holds.
entry_state = ((ema_fast > ema_slow) & (rsi14 > 50)).astype(bool)
prev_entry_state = entry_state.shift(1, fill_value=False).astype(bool)
raw_entries = entry_state & ~prev_entry_state

# Exit fires on the bar the 10-day EMA crosses below the 30-day EMA.
raw_exits = (ema_fast < ema_slow) & (ema_fast.shift(1) >= ema_slow.shift(1))

# Shift by 1 bar so a signal evaluated on day T executes at day T+1's open.
entries = raw_entries.vbt.signals.fshift(1)
exits = raw_exits.vbt.signals.fshift(1)

# This strategy does not enforce a hard cash constraint: a position can be
# opened using unrealized gains elsewhere in the portfolio, not just literal
# starting cash (the same behavior as the R/blotter implementation). The
# backtest is run with a large nominal cash buffer so VectorBT's own
# cash-sharing check never artificially rejects a fill, then every reported
# figure below is rebased to the real $100,000 starting capital.
init_cash_sim = 10_000_000

portfolio = vbt.Portfolio.from_signals(
    open_price,
    entries=entries,
    exits=exits,
    price=open_price,
    init_cash=init_cash_sim,
    cash_sharing=True,
    group_by=True,
    size=300,            # Fixed 300-share order per entry
    size_type='amount',
    fees=0.0,
    slippage=0.0005,      # 5 bps execution friction
    freq='1D'
)

# Rebase the simulated equity curve onto the real starting capital.
strategy_equity = portfolio.value() - init_cash_sim + init_cash
strategy_ret = strategy_equity.pct_change().dropna()

final_val = strategy_equity.iloc[-1]
total_ret = (final_val / init_cash - 1) * 100
sharpe = strategy_ret.mean() / strategy_ret.std() * np.sqrt(252)
maxdd = -((strategy_equity / strategy_equity.cummax()) - 1).min() * 100

trade_stats = portfolio.trades.records_readable
n_trades = int((trade_stats['Status'] == 'Closed').sum())
closed = trade_stats[trade_stats['Status'] == 'Closed']
win_rate = (closed['PnL'] > 0).mean() * 100
profit_factor = closed.loc[closed['PnL'] > 0, 'PnL'].sum() / abs(closed.loc[closed['PnL'] < 0, 'PnL'].sum())

per_symbol = closed.groupby('Column').agg(
    Num_Trades=('PnL', 'count'),
    Percent_Positive=('PnL', lambda s: (s > 0).mean() * 100),
    Net_PnL=('PnL', 'sum'),
    Avg_PnL=('PnL', 'mean'),
).reset_index().rename(columns={'Column': 'Symbol'})