# Sector Momentum Breakout & Rotation -- Python / VectorBT
# Analysis & validation code: benchmarks, risk diagnostics, cash overlay,
# parameter optimization, walk-forward validation.
# Run sector-momentum-python-strategy.py first, in the same Python session.
# Full notebook: https://quantstr.at/wp-content/uploads/reports/sector-momentum-python.html

spy_close = vbt.YFData.download("SPY", start=start_date, end=end_date).get('Close')
spy_close = spy_close[spy_close.index.date >= start_calendar_date]
spy_equity = init_cash * spy_close / spy_close.iloc[0]

basket_equity = ((init_cash / len(symbols)) * close / close.iloc[0]).sum(axis=1)

comparison = pd.DataFrame({
    "Strategy": strategy_equity,
    "SPY": spy_equity,
    "Sector_Basket": basket_equity,
}).dropna()

fig, ax = plt.subplots(figsize=(9, 5))
comparison.plot(ax=ax, color=["#00E08F", "#F59E0B", "#2563EB"], linewidth=2)
ax.set_title("Strategy vs. Benchmark Equity Curves")
ax.set_ylabel("Equity ($)")
plt.show()

sim_cash = portfolio.cash()
total_sim_value = portfolio.value()
exposure_pct = 100 * (1 - sim_cash / total_sim_value)

fig, ax = plt.subplots(figsize=(9, 5))
exposure_pct.plot(ax=ax, color="#00E08F", linewidth=1.5)
ax.axhline(0, linestyle="--", color="gray")
ax.set_title("Portfolio Exposure Over Time")
ax.set_ylabel("% of Equity Invested")
plt.show()
print(f"Average invested exposure: {exposure_pct.mean():.1f}%")

bh_returns = (close.iloc[-1] / close.iloc[0] - 1) * 100

# BIL = 1-3 month T-Bill ETF (near-zero duration); BND = intermediate-duration
# bonds; TLT = long-duration Treasuries.
cash_real = sim_cash - init_cash_sim + init_cash

def bond_overlay(bond_symbol):
    bond_close = vbt.YFData.download(bond_symbol, start=start_date, end=end_date).get('Close')
    bond_close = bond_close[bond_close.index.date >= start_calendar_date]
    bond_ret = bond_close.pct_change().reindex(cash_real.index).fillna(0)
    cash_flow = cash_real.diff()
    cash_flow.iloc[0] = 0
    cash_bond = cash_real.copy()
    for i in range(1, len(cash_bond)):
        cash_bond.iloc[i] = cash_bond.iloc[i - 1] * (1 + bond_ret.iloc[i]) + cash_flow.iloc[i]
    new_equity = cash_bond + (strategy_equity - cash_real)
    new_ret = new_equity.pct_change().dropna()
    total_return_ = (new_equity.iloc[-1] / init_cash - 1) * 100
    sharpe_ = new_ret.mean() / new_ret.std() * (252 ** 0.5)
    max_dd_ = -((new_equity / new_equity.cummax()) - 1).min() * 100
    return total_return_, sharpe_, max_dd_

overlay_rows = [{"Overlay": "Baseline", "Total Return %": total_ret, "Sharpe": sharpe, "Max DD %": maxdd}]
for sym in ["BIL", "BND", "TLT"]:
    tr, sh, dd = bond_overlay(sym)
    overlay_rows.append({"Overlay": sym, "Total Return %": tr, "Sharpe": sh, "Max DD %": dd})
overlay_df = pd.DataFrame(overlay_rows).round(2)

rolling_sharpe = strategy_ret.rolling(63).apply(lambda x: x.mean() / x.std() * np.sqrt(252), raw=True)

fig, ax = plt.subplots(figsize=(9, 5))
rolling_sharpe.plot(ax=ax, color="#00E08F")
ax.axhline(0, linestyle="--", color="gray")
ax.set_title("Rolling 3-Month Annualized Sharpe Ratio")
plt.show()

drawdown = (strategy_equity / strategy_equity.cummax() - 1) * 100
dd = drawdown.copy()
troughs, in_dd = [], False
for date, val in dd.items():
    if val < 0 and not in_dd:
        in_dd, start, trough_val, trough_date = True, date, val, date
    elif val < 0 and in_dd:
        if val < trough_val:
            trough_val, trough_date = val, date
    elif val >= 0 and in_dd:
        in_dd = False
        troughs.append({"From": start.date(), "Trough": trough_date.date(), "To": date.date(), "Depth %": trough_val})
if in_dd:
    troughs.append({"From": start.date(), "Trough": trough_date.date(), "To": None, "Depth %": trough_val})
top5 = pd.DataFrame(troughs).sort_values("Depth %").head(5).reset_index(drop=True)
top5["Depth %"] = top5["Depth %"].round(2)

fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(close.index, close["XLK"], color="#1a1714", linewidth=1, label="XLK Close")
ax.plot(close.index, ema_fast["XLK"], color="#2563EB", linewidth=1, label="10-day EMA")
ax.plot(close.index, ema_slow["XLK"], color="#D9603E", linewidth=1, label="30-day EMA")
buy_dates = close.index[entries["XLK"].fillna(False)]
sell_dates = close.index[exits["XLK"].fillna(False)]
ax.scatter(buy_dates, close.loc[buy_dates, "XLK"], marker="^", color="#00E08F", s=60, zorder=5, label="Entry")
ax.scatter(sell_dates, close.loc[sell_dates, "XLK"], marker="v", color="#B3392F", s=60, zorder=5, label="Exit")
ax.set_title("XLK Price, 10/30 EMA & Signals")
ax.legend(loc="upper left", fontsize=8)
plt.show()

fast_grid = [5, 10, 15, 20]
slow_grid = [20, 30, 40, 50]
combos = [(f, s) for f in fast_grid for s in slow_grid if f < s]

def compute_signals(fast_n, slow_n, px):
    ema_f = px.apply(lambda col: ema(col, fast_n))
    ema_s = px.apply(lambda col: ema(col, slow_n))
    rsi_ = px.apply(lambda col: rsi(col, 14))
    state = ((ema_f > ema_s) & (rsi_ > 50)).astype(bool)
    prev = state.shift(1, fill_value=False).astype(bool)
    e = (state & ~prev).vbt.signals.fshift(1)
    x = ((ema_f < ema_s) & (ema_f.shift(1) >= ema_s.shift(1))).vbt.signals.fshift(1)
    return e, x

def run_window(fast_n, slow_n, win_start, win_end, cash_in, px, op):
    e_full, x_full = compute_signals(fast_n, slow_n, px)
    e, x = e_full.loc[win_start:win_end], x_full.loc[win_start:win_end]
    o = op.loc[win_start:win_end]
    return vbt.Portfolio.from_signals(o, entries=e, exits=x, price=o, init_cash=cash_in,
                                       cash_sharing=True, group_by=True, size=300,
                                       size_type='amount', fees=0.0, slippage=0.0005, freq='1D')

opt_rows = []
for fast_n, slow_n in combos:
    pf = run_window(fast_n, slow_n, start_date, end_date, init_cash_sim, close, open_price)
    opt_rows.append({"fast_n": fast_n, "slow_n": slow_n, "net_pl": pf.final_value() - init_cash_sim})
opt_df = pd.DataFrame(opt_rows).sort_values("net_pl", ascending=False).reset_index(drop=True)
opt_df["net_pl"] = opt_df["net_pl"].round(0)

wf_end_date = pd.Timestamp.today().strftime("%Y-%m-%d")
wf_data = vbt.YFData.download(symbols, start=start_date, end=wf_end_date)
wf_close = wf_data.get('Close')
wf_open = wf_data.get('Open')
wf_close = wf_close[wf_close.index.date >= start_calendar_date]
wf_open = wf_open[wf_open.index.date >= start_calendar_date]

month_end_mask = wf_close.index.to_series().groupby(
    [wf_close.index.year, wf_close.index.month]).transform('idxmax') == wf_close.index.to_series()
me_dates = wf_close.index[month_end_mask.values]

k_train, k_test = 18, 6
windows = []
w = 0
while (k_train + w * k_test) < len(me_dates):
    tr_end_i = k_train - 1 + w * k_test
    te_end_i = k_train - 1 + (w + 1) * k_test
    tr_start = wf_close.index[0] if w == 0 else me_dates[w * k_test - 1] + pd.Timedelta(days=1)
    tr_end = me_dates[tr_end_i]
    te_start = tr_end + pd.Timedelta(days=1)
    te_end = me_dates[te_end_i] if te_end_i < len(me_dates) else wf_close.index[-1]
    if te_start > wf_close.index[-1]:
        break
    windows.append((tr_start, tr_end, te_start, te_end))
    w += 1

windows_df = pd.DataFrame(windows, columns=["training_start", "training_end", "testing_start", "testing_end"])

chosen_rows = []
oos_equity_parts = []
running_cash = init_cash
for (tr_start, tr_end, te_start, te_end) in windows:
    best_combo, best_obj = None, -np.inf
    for fast_n, slow_n in combos:
        pf_train = run_window(fast_n, slow_n, tr_start, tr_end, init_cash_sim, wf_close, wf_open)
        obj = pf_train.final_value() - init_cash_sim
        if obj > best_obj:
            best_obj, best_combo = obj, (fast_n, slow_n)
    chosen_rows.append({"training_start": tr_start.date(), "training_end": tr_end.date(),
                         "fast_n": best_combo[0], "slow_n": best_combo[1], "train_pl": round(best_obj, 0)})

    pf_test = run_window(best_combo[0], best_combo[1], te_start, te_end, init_cash_sim, wf_close, wf_open)
    window_equity = pf_test.value() - init_cash_sim + running_cash
    oos_equity_parts.append(window_equity)
    running_cash = float(window_equity.iloc[-1])

chosen_df = pd.DataFrame(chosen_rows)

oos_equity = pd.concat(oos_equity_parts)
oos_equity = oos_equity[~oos_equity.index.duplicated(keep="last")]
live_start, oos_end = oos_equity.index[0], oos_equity.index[-1]

wf_total_ret = (oos_equity.iloc[-1] / init_cash - 1) * 100
oos_ret = oos_equity.pct_change().dropna()
wf_sharpe = oos_ret.mean() / oos_ret.std() * np.sqrt(252)
wf_maxdd = -((oos_equity / oos_equity.cummax()) - 1).min() * 100

print(f"Walk-forward out-of-sample: ${oos_equity.iloc[-1]:,.0f} "
      f"({wf_total_ret:.1f}%), Sharpe {wf_sharpe:.2f}, Max DD {wf_maxdd:.1f}% "
      f"over {live_start.date()} to {oos_end.date()}")

def perf_stats(series):
    s = series.dropna()
    r = s.pct_change().dropna()
    return pd.Series({
        "Total_Return_Pct": (s.iloc[-1] / s.iloc[0] - 1) * 100,
        "Ann_Sharpe": r.mean() / r.std() * np.sqrt(252),
        "Max_Drawdown_Pct": -((s / s.cummax()) - 1).min() * 100
    })

fixed_pf = run_window(10, 30, live_start, oos_end, init_cash_sim, wf_close, wf_open)
fixed_equity = fixed_pf.value() - init_cash_sim + init_cash

hind_rows = []
for fast_n, slow_n in combos:
    pf = run_window(fast_n, slow_n, live_start, oos_end, init_cash_sim, wf_close, wf_open)
    hind_rows.append((fast_n, slow_n, pf.final_value()))
hb = max(hind_rows, key=lambda r: r[2])
hind_pf = run_window(hb[0], hb[1], live_start, oos_end, init_cash_sim, wf_close, wf_open)
hind_equity = hind_pf.value() - init_cash_sim + init_cash

spy_wf = vbt.YFData.download("SPY", start=start_date, end=wf_end_date).get('Close')
spy_wf = spy_wf[spy_wf.index.date >= start_calendar_date].loc[live_start:oos_end]
basket_wf = ((wf_close.loc[live_start:oos_end]) / wf_close.loc[live_start]).mean(axis=1) * init_cash

wf_comparison = pd.DataFrame({
    "Walk-forward (out-of-sample)": perf_stats(oos_equity),
    "Fixed EMA 10/30 (baseline rule)": perf_stats(fixed_equity),
    f"Hindsight-best fixed {hb[0]}/{hb[1]} (unattainable live)": perf_stats(hind_equity),
    "SPY buy & hold": perf_stats(spy_wf),
    "Equal-weight sector basket": perf_stats(basket_wf),
}).T.round(2)

def rebase(s):
    return s / s.iloc[0] * 100 - 100

fig, ax = plt.subplots(figsize=(9, 5.5))
rebase(spy_wf).plot(ax=ax, color="#F59E0B", linewidth=2, label="SPY Buy & Hold")
rebase(basket_wf).plot(ax=ax, color="#2563EB", linewidth=2, label="Sector Basket Buy & Hold")
rebase(fixed_equity).plot(ax=ax, color="gray", linewidth=2, linestyle="--", label="Fixed EMA 10/30 (baseline rule)")
rebase(oos_equity).plot(ax=ax, color="#00E08F", linewidth=3, label="Walk-forward (params re-chosen each window)")
for _, row in chosen_df.iterrows():
    ax.axvline(pd.Timestamp(row["training_end"]) + pd.Timedelta(days=1), color="gray", linestyle=":", alpha=0.5)
ax.set_title("Walk-Forward Out-of-Sample Return vs. Benchmarks (same start date)")
ax.set_ylabel("Cumulative Return (%)")
ax.legend(loc="upper left", fontsize=8)
plt.figtext(0.5, -0.02, "Dotted vertical lines = re-optimization dates (start of each test window)",
            ha="center", fontsize=8, color="gray")
plt.show()
