# Keltner Channel Volatility Breakout with ADX Trend Filter -- Python / VectorBT
#
# Layout convention (read by the Quarto report generator): "# %% <title>" starts
# a section; "#'" lines are prose shown above the section's code.

# %% Environment
import numpy as np
import pandas as pd
import vectorbt as vbt

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

# %% Universe & Parameters
#' Same universe, window and sizing as the R report: one equal dollar slot per
#' symbol, long-only, no leverage, no adding to an open position.
symbols = ["QQQ", "NVDA", "AAPL", "MSFT", "AMD", "AMZN"]
start_date = "2023-08-29"
end_date = "2026-08-27"
init_equity = 100_000

ema_n = 20       # Keltner centre line: EMA of the close
atr_n = 10       # Keltner band width: ATR period
atr_mult = 2.0   # Keltner band width: ATR multiplier
adx_n = 14       # trend-strength filter period
adx_min = 25     # entries require ADX above this level
slot_usd = init_equity / len(symbols)
cost_rate = 0.0005  # 5 bps of traded value on every fill

# %% Data
#' Daily bars from Yahoo Finance via `yfinance`, with OHLC adjusted for splits
#' and dividends.
data = vbt.YFData.download(symbols, start=start_date, end=end_date, missing_columns="nan")
close = data.get("Close")[symbols]
open_ = data.get("Open")[symbols]
high = data.get("High")[symbols]
low = data.get("Low")[symbols]

# Keep the same trading-day range as the R report (calendar-date filter).
in_range = close.index.date >= pd.Timestamp(start_date).date()
close, open_, high, low = close[in_range], open_[in_range], high[in_range], low[in_range]

worst = close.pct_change().abs().max()
assert (worst < 0.40).all(), f"suspect unadjusted split: {worst[worst >= 0.40].to_dict()}"
print("Bars downloaded per symbol:")
print(close.count())

# %% Indicators
#' EMA uses alpha = 2/(n+1) and Wilder smoothing uses alpha = 1/n; both are
#' SMA-seeded, matching the R report's `TTR::EMA()`, `TTR::ATR()` and `TTR::ADX()`.
def sma_seeded(df: pd.DataFrame, n: int, alpha: float) -> pd.DataFrame:
    x = df.to_numpy(dtype=float)
    out = np.full(x.shape, np.nan)
    for j in range(x.shape[1]):
        col = x[:, j]
        valid = np.flatnonzero(~np.isnan(col))
        if len(valid) < n:
            continue
        s = valid[0] + n - 1
        out[s, j] = col[valid[0]:s + 1].mean()
        for i in range(s + 1, len(col)):
            out[i, j] = alpha * col[i] + (1 - alpha) * out[i - 1, j]
    return pd.DataFrame(out, index=df.index, columns=df.columns)

ema = lambda df, n: sma_seeded(df, n, 2.0 / (n + 1))
wilder = lambda df, n: sma_seeded(df, n, 1.0 / n)

prev_close = close.shift(1)
true_range = pd.DataFrame(
    np.maximum.reduce([(high - low).to_numpy(), (high - prev_close).abs().to_numpy(), (low - prev_close).abs().to_numpy()]),
    index=close.index, columns=close.columns)
true_range.iloc[0] = (high - low).iloc[0]

mid = ema(close, ema_n)
atr = wilder(true_range, atr_n)
upper = mid + atr_mult * atr
lower = mid - atr_mult * atr

up_move = high.diff()
down_move = -low.diff()
plus_dm = up_move.where((up_move > down_move) & (up_move > 0), 0.0)
minus_dm = down_move.where((down_move > up_move) & (down_move > 0), 0.0)
tr_adx = wilder(true_range, adx_n)
di_pos = 100 * wilder(plus_dm, adx_n) / tr_adx
di_neg = 100 * wilder(minus_dm, adx_n) / tr_adx
dx = 100 * (di_pos - di_neg).abs() / (di_pos + di_neg)
adx = wilder(dx.dropna(how="all"), adx_n).reindex(close.index)

# %% Signals
#' Entry fires on the first bar where the close is above the upper band AND ADX is
#' above its threshold. Exit fires when the close crosses below the EMA centre line.
#' A signal evaluated on day T is executed at day T+1's open.
entry_state = ((close > upper) & (adx > adx_min)).fillna(False).astype(bool)
raw_entries = entry_state & ~entry_state.shift(1, fill_value=False)
raw_exits = ((close < mid) & (close.shift(1) >= mid.shift(1))).fillna(False).astype(bool)

entries = raw_entries.shift(1, fill_value=False).astype(bool)
exits = raw_exits.shift(1, fill_value=False).astype(bool)

# %% Run the Backtest
#' Each entry buys one dollar slot at the next open; `accumulate=False` ignores a
#' second entry while a position is open. Every fill pays 5 bps of traded value.
portfolio = vbt.Portfolio.from_signals(
    close,
    entries=entries,
    exits=exits,
    price=open_,
    size=slot_usd,
    size_type="value",
    fees=cost_rate,
    init_cash=init_equity,
    cash_sharing=True,
    group_by=True,
    accumulate=False,
    freq="1D",
)

# %% Headline Performance
equity = portfolio.value()
returns = equity.pct_change().dropna()
final_val = float(equity.iloc[-1])
total_ret = (final_val / init_equity - 1) * 100
sharpe = returns.mean() / returns.std() * np.sqrt(252)
max_dd = -((equity / equity.cummax()) - 1).min() * 100

trades = portfolio.trades.records_readable
closed = trades[trades["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())

print(pd.Series({
    "Final Equity": f"${final_val:,.0f}",
    "Total Return": f"{total_ret:.2f}%",
    "Annualized Sharpe": f"{sharpe:.2f}",
    "Max Drawdown": f"{max_dd:.2f}%",
    "Closed Trades": int(len(closed)),
    "Win Rate": f"{win_rate:.1f}%",
    "Profit Factor": f"{profit_factor:.2f}",
}).to_string())

# %% Trade Statistics
#' One row per symbol: closed-trade count, net P&L, win rate and profit factor.
per_symbol = closed.groupby("Column").agg(
    Num_Trades=("PnL", "count"),
    Net_PnL=("PnL", "sum"),
    Percent_Positive=("PnL", lambda s: (s > 0).mean() * 100),
    Avg_PnL=("PnL", "mean"),
).reset_index().rename(columns={"Column": "Symbol"})
print(per_symbol.round(2).to_string(index=False))