!pip install pandas numpy plotly

import warnings
warnings.filterwarnings("ignore")

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from IPython.display import display
import plotly.io as pio

def generate_data(periods: int, volatility_scale: float = 0.0012, wick_scale: float = 0.0006,
                   seed: int = None) -> pd.DataFrame:
    """
    Generate synthetic OHLCV data via a simple random walk.

    volatility_scale : std of close-to-close returns per bar, as a fraction of price.
                        ~0.0012 (0.12%) approximates realistic 1-minute crypto volatility.
    wick_scale        : additional std for high/low wicks beyond the open/close body.
    """
    if seed is not None:
        np.random.seed(seed)

    start_date     = pd.to_datetime("2024-01-01 00:00:00+00:00")
    datetime_index = pd.date_range(start_date, periods=periods, freq="1min", tz="UTC")
    price_data = []; last_close = 42000.0

    for _ in range(periods):
        open_price  = last_close + np.random.normal(0, last_close * volatility_scale * 0.1)
        close_price = open_price + np.random.normal(0, last_close * volatility_scale)
        body_high   = max(open_price, close_price)
        body_low    = min(open_price, close_price)
        high_price  = max(body_high + abs(np.random.normal(0, last_close * wick_scale)),
                          open_price, close_price)
        low_price   = min(body_low  - abs(np.random.normal(0, last_close * wick_scale)),
                          open_price, close_price)
        if high_price < low_price:
            high_price, low_price = low_price, high_price
        price_data.append({
            "open":  max(1, int(open_price)),
            "high":  max(1, int(high_price)),
            "low":   max(1, int(low_price)),
            "close": max(1, int(close_price)),
        })
        last_close = close_price

    df = pd.DataFrame(price_data, index=datetime_index)
    df.index.name = "datetime"
    df["volume"]   = np.random.uniform(100.0, 500.0, periods)
    df["datetime"] = df.index.to_series()
    return df.reset_index(drop=True)

df = generate_data(500, seed=42)
display(df.head())
print(f"1-minute return std: {df['close'].pct_change().std()*100:.3f}%  "
      f"(realistic range for crypto is roughly 0.05%-0.15%)")

def apply_slippage(
    price:             float,
    direction:         int,
    model:             str   = "fixed",
    slippage_pct:      float = 0.0003,
    atr:               float = None,
    atr_multiplier:    float = 0.1,
    avg_volume:        float = None,
    order_size:        float = None,
    volume_impact_k:   float = 0.001,
) -> float:
    """
    Compute the fill price after applying the specified slippage model.

    Parameters
    ----------
    direction        : +1 for buy (fills higher), -1 for sell (fills lower).
    model            : One of 'fixed', 'random', 'atr_scaled', 'volume_impact'.
    slippage_pct     : Base slippage percentage for 'fixed' and 'random' models.
    atr              : Current ATR value, in price units (required for 'atr_scaled').
    atr_multiplier   : Fraction of ATR applied as slippage (e.g., 0.1 = 10% of ATR).
    avg_volume       : Rolling average traded volume, in asset units (required for 'volume_impact').
    order_size       : Size of the order in quote currency / dollars (required for 'volume_impact').
    volume_impact_k  : Scaling coefficient for the volume-impact model.

    Returns
    -------
    float : Adjusted fill price after slippage.
    """
    if model == "fixed":
        slip = slippage_pct

    elif model == "random":
        slip = np.random.uniform(0, slippage_pct)

    elif model == "atr_scaled":
        if atr is None or np.isnan(atr) or price == 0:
            slip = slippage_pct
        else:
            slip = (atr * atr_multiplier) / price

    elif model == "volume_impact":
        if avg_volume is None or avg_volume == 0 or order_size is None or np.isnan(avg_volume):
            slip = slippage_pct
        else:
            dollar_volume = avg_volume * price          # FIX: convert units volume -> dollar volume
            slip = volume_impact_k * (order_size / dollar_volume)

    else:
        raise ValueError(f"Unknown slippage model: {model}")

    # Direction determines sign: buys fill higher, sells fill lower
    return price * (1 + direction * slip)

def compute_atr(df: pd.DataFrame, window: int = 14) -> pd.Series:
    """Simple-moving-average True Range (a lightweight stand-in for Wilder's
    smoothed ATR — close enough for illustrating volatility-scaled slippage,
    but note it is not the textbook Wilder ATR)."""
    tr = pd.concat([
        df["high"] - df["low"],
        (df["high"] - df["close"].shift(1)).abs(),
        (df["low"]  - df["close"].shift(1)).abs(),
    ], axis=1).max(axis=1)
    return tr.rolling(window).mean()

def backtest_slippage_comparison(
    df:              pd.DataFrame,
    initial_capital: float = 10_000.0,
    fee_pct:         float = 0.0005,
    slippage_pct:    float = 0.0003,
) -> tuple[pd.DataFrame, dict]:
    df = df.copy().sort_values("datetime", ignore_index=True)
    df["fast_ma"]    = df["close"].rolling(10).mean()
    df["slow_ma"]    = df["close"].rolling(30).mean()
    df["atr"]        = compute_atr(df, window=14)
    df["avg_volume"] = df["volume"].rolling(20).mean()

    results   = {}
    n_trades  = {}
    order_size_fraction = 0.05  # 5% of capital per entry, used for volume-impact sizing

    for model in ["none", "fixed", "random", "atr_scaled", "volume_impact"]:
        cash = initial_capital
        pos  = 0.0
        equity_curve = []
        trades = 0

        for _, row in df.iterrows():
            if pd.isna(row["fast_ma"]) or pd.isna(row["slow_ma"]):
                equity_curve.append(cash + pos * row["close"])
                continue

            sig   = 1 if row["fast_ma"] > row["slow_ma"] else 0
            price = row["close"]

            if sig == 1 and pos == 0:
                fill = price if model == "none" else apply_slippage(
                    price=price, direction=+1, model=model,
                    slippage_pct=slippage_pct,
                    atr=row["atr"], atr_multiplier=0.1,
                    avg_volume=row["avg_volume"], order_size=initial_capital * order_size_fraction,
                )
                units = (cash * 0.95) / fill
                cash -= units * fill * (1 + fee_pct)
                pos    = units
                trades += 1

            elif sig == 0 and pos > 0:
                fill = price if model == "none" else apply_slippage(
                    price=price, direction=-1, model=model,
                    slippage_pct=slippage_pct,
                    atr=row["atr"], atr_multiplier=0.1,
                    avg_volume=row["avg_volume"], order_size=initial_capital * order_size_fraction,
                )
                cash += pos * fill * (1 - fee_pct)
                pos    = 0.0
                trades += 1

            equity_curve.append(cash + pos * price)

        results[f"equity_{model}"] = equity_curve
        n_trades[model] = trades

    for col, vals in results.items():
        df[col] = vals

    return df, n_trades

df_slip, n_trades = backtest_slippage_comparison(df)

print("Trade counts (should be identical across models if only fill price varies):")
print(n_trades)
print()

summary_rows = []
for model in ["none", "fixed", "random", "atr_scaled", "volume_impact"]:
    col   = f"equity_{model}"
    final = df_slip[col].iloc[-1]
    ret   = (final / 10_000 - 1) * 100
    summary_rows.append((model, final, ret))
    print(f"{model:14s}: Final ${final:,.2f}  ({ret:+.2f}%)")

fig = make_subplots(rows=1, cols=1)

colors = {'none': '#888888', 'fixed': '#E4572E', 'random': '#F3A712',
          'atr_scaled': '#2E86AB', 'volume_impact': '#3DA35D'}

for model in ['none', 'fixed', 'random', 'atr_scaled', 'volume_impact']:
    col = f"equity_{model}"
    fig.add_trace(go.Scatter(
        x=df_slip["datetime"],
        y=df_slip[col],
        mode='lines',
        name=f'{model.replace("_", " ").title()} Slippage',
        line=dict(color=colors[model])
    ), row=1, col=1)

fig.update_layout(
    title_text='<b>Slippage Model Comparison — Equity Curve Impact</b>',
    height=550,
    xaxis_rangeslider_visible=False,
    hovermode='x unified',
    template='plotly_white',
)

fig.update_yaxes(title_text='Equity ($)', row=1, col=1)
fig.update_xaxes(title_text='Time', row=1, col=1)

fig.show()

baseline = df_slip["equity_none"].iloc[-1]

cost_rows = []
for model in ["fixed", "random", "atr_scaled", "volume_impact"]:
    final = df_slip[f"equity_{model}"].iloc[-1]
    dollar_cost = baseline - final
    pct_cost    = dollar_cost / baseline * 100
    cost_rows.append((model, dollar_cost, pct_cost))

cost_df = pd.DataFrame(cost_rows, columns=["model", "dollar_cost", "pct_cost"])
display(cost_df)

fig2 = go.Figure(go.Bar(
    x=cost_df["model"].str.replace("_", " ").str.title(),
    y=cost_df["pct_cost"],
    marker_color=[colors[m] for m in cost_df["model"]],
    text=[f"{v:.2f}%" for v in cost_df["pct_cost"]],
    textposition="outside",
))
fig2.update_layout(
    title_text="<b>Slippage Cost vs. No-Slippage Baseline</b>",
    yaxis_title="Cost (% of baseline final equity)",
    template="plotly_white",
    height=450,
)
fig2.show()