!pip install numpy pandas matplotlib seaborn --quiet

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from dataclasses import dataclass, field
from typing import List, Dict, Optional
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline
plt.rcParams['figure.figsize'] = (14, 5)
plt.rcParams['axes.spines.top']   = False
plt.rcParams['axes.spines.right'] = False
print('Imports ready.')

# --- Configuration ---
GRID_CENTER      = 50_000.0   # central price (approximate current price)
GRID_RANGE_PCT   = 0.20       # ±20% range around centre
N_LEVELS         = 20         # number of grid levels
ORDER_SIZE_USD   = 500.0      # USD allocated per grid level
SIMULATION_DAYS  = 90         # days to simulate
GRID_TYPE        = 'geometric' # 'arithmetic' or 'geometric'
print('Config ready.')

@dataclass
class GridLevel:
    """
    A single level in the trading grid.

    Attributes
    ----------
    price : float  Price at which a buy or sell order sits.
    order_size_btc : float  BTC quantity per fill.
    has_buy_order : bool  Whether a buy order is active at this level.
    has_sell_order : bool  Whether a sell order is active at this level.
    buy_fills : int  Cumulative buy fills at this level.
    sell_fills : int  Cumulative sell fills at this level.
    """
    price: float
    order_size_btc: float
    has_buy_order: bool = True
    has_sell_order: bool = False
    buy_fills: int = 0
    sell_fills: int = 0


def create_arithmetic_grid(
    center: float,
    range_pct: float,
    n_levels: int,
    order_size_usd: float
) -> List[GridLevel]:
    """
    Create a grid with equal dollar spacing between levels.

    Parameters
    ----------
    center : float  Reference price.
    range_pct : float  Half-range as fraction of center (e.g. 0.20 = ±20%).
    n_levels : int  Total number of levels.
    order_size_usd : float  USD value per order.

    Returns
    -------
    List[GridLevel]  Sorted ascending list of grid levels.
    """
    lo = center * (1 - range_pct)
    hi = center * (1 + range_pct)
    prices = np.linspace(lo, hi, n_levels)
    return [GridLevel(price=p, order_size_btc=order_size_usd / p) for p in prices]


def create_geometric_grid(
    center: float,
    range_pct: float,
    n_levels: int,
    order_size_usd: float
) -> List[GridLevel]:
    """
    Create a grid with equal percentage spacing between levels.

    Parameters
    ----------
    center : float  Reference price.
    range_pct : float  Half-range as fraction of center.
    n_levels : int  Total number of levels.
    order_size_usd : float  USD value per order.

    Returns
    -------
    List[GridLevel]  Sorted ascending list of grid levels.

    Notes
    -----
    Geometric spacing is preferred for crypto because percentage moves
    are more natural than dollar moves at high price levels.
    """
    lo = center * (1 - range_pct)
    hi = center * (1 + range_pct)
    prices = np.geomspace(lo, hi, n_levels)
    return [GridLevel(price=p, order_size_btc=order_size_usd / p) for p in prices]


grid = (create_geometric_grid if GRID_TYPE == 'geometric' else create_arithmetic_grid)(
    GRID_CENTER, GRID_RANGE_PCT, N_LEVELS, ORDER_SIZE_USD
)
print(f'Grid type: {GRID_TYPE}')
print(f'Range: ${grid[0].price:,.0f} — ${grid[-1].price:,.0f}')
print(f'Levels: {len(grid)}')
print(f'Spacing: {((grid[1].price / grid[0].price - 1) * 100):.2f}% (between levels 0 and 1)')

def generate_synthetic_price_series(
    start: float,
    n_days: int,
    annual_vol: float = 0.60,
    drift: float = 0.0,
    seed: int = 42
) -> np.ndarray:
    """
    Generate a synthetic BTC price series using geometric Brownian motion.

    Parameters
    ----------
    start : float  Starting price.
    n_days : int  Number of daily prices.
    annual_vol : float  Annual volatility (e.g. 0.60 = 60%).
    drift : float  Annual drift term.
    seed : int  Random seed.

    Returns
    -------
    np.ndarray  Daily close prices.
    """
    rng = np.random.default_rng(seed)
    dt  = 1 / 365
    rets = (drift - 0.5 * annual_vol**2) * dt + annual_vol * np.sqrt(dt) * rng.standard_normal(n_days)
    return start * np.exp(np.cumsum(rets))


def simulate_grid_trading(
    grid: List[GridLevel],
    prices: np.ndarray
) -> pd.DataFrame:
    """
    Simulate grid trading by checking price crossings each bar.

    When price crosses a level downward, the buy order at that level fills
    and a sell order is placed one level above. When price crosses upward,
    the sell order fills and a buy order is placed one level below.

    Parameters
    ----------
    grid : List[GridLevel]  Initialized grid levels.
    prices : np.ndarray  Price series (daily or intraday).

    Returns
    -------
    pd.DataFrame
        Trade log with columns: bar, price, side, level_idx,
        fill_price, qty_btc, pnl_usd.
    """
    import copy
    grid = copy.deepcopy(grid)
    # Initialize: set buy orders below current price, sell orders above
    current = prices[0]
    for i, lvl in enumerate(grid):
        lvl.has_buy_order  = (lvl.price < current)
        lvl.has_sell_order = (lvl.price > current)

    trades = []
    prev_price = prices[0]
    cash = 0.0
    btc  = 0.0

    for bar, price in enumerate(prices):
        lo_p = min(price, prev_price)
        hi_p = max(price, prev_price)

        for i, lvl in enumerate(grid):
            # Buy fill: price moved down through this level
            if lvl.has_buy_order and lo_p <= lvl.price <= hi_p:
                cost = lvl.order_size_btc * lvl.price
                cash -= cost
                btc  += lvl.order_size_btc
                lvl.buy_fills += 1
                lvl.has_buy_order  = False
                if i + 1 < len(grid):
                    grid[i + 1].has_sell_order = True  # place sell one level up
                trades.append({'bar': bar, 'price': price, 'side': 'BUY',
                                'fill_price': lvl.price, 'qty_btc': lvl.order_size_btc,
                                'pnl_usd': 0.0})

            # Sell fill: price moved up through this level
            elif lvl.has_sell_order and lo_p <= lvl.price <= hi_p:
                revenue = lvl.order_size_btc * lvl.price
                cash   += revenue
                btc    -= lvl.order_size_btc
                lvl.sell_fills += 1
                lvl.has_sell_order = False
                if i - 1 >= 0:
                    grid[i - 1].has_buy_order = True  # place buy one level down
                # Grid profit = sell_price - buy_price (one level below)
                buy_price = grid[i-1].price if i > 0 else lvl.price
                grid_pnl  = (lvl.price - buy_price) * lvl.order_size_btc
                trades.append({'bar': bar, 'price': price, 'side': 'SELL',
                                'fill_price': lvl.price, 'qty_btc': lvl.order_size_btc,
                                'pnl_usd': grid_pnl})

        prev_price = price

    return pd.DataFrame(trades) if trades else pd.DataFrame(columns=['bar','price','side','fill_price','qty_btc','pnl_usd'])


prices = generate_synthetic_price_series(GRID_CENTER, SIMULATION_DAYS)
trades_df = simulate_grid_trading(grid, prices)
print(f'Total trades: {len(trades_df)}')
print(f'Buy fills: {(trades_df["side"]=="BUY").sum()}, Sell fills: {(trades_df["side"]=="SELL").sum()}')
print(f'Total grid PnL: ${trades_df["pnl_usd"].sum():,.2f}')

def compute_grid_performance(
    trades_df: pd.DataFrame,
    prices: np.ndarray,
    order_size_usd: float,
    n_levels: int
) -> dict:
    """
    Summarize grid trading performance metrics.

    Parameters
    ----------
    trades_df : pd.DataFrame  Trade log from simulate_grid_trading().
    prices : np.ndarray  Price series used in simulation.
    order_size_usd : float  USD per grid order.
    n_levels : int  Number of grid levels.

    Returns
    -------
    dict  Performance metrics.
    """
    total_capital = order_size_usd * n_levels
    grid_pnl      = trades_df['pnl_usd'].sum()
    sell_trades   = trades_df[trades_df['side'] == 'SELL']
    n_completed   = len(sell_trades)  # completed pairs
    price_change  = (prices[-1] / prices[0] - 1) * 100
    bh_pnl        = (prices[-1] - prices[0]) * (total_capital / prices[0])

    return {
        'total_capital_usd':   total_capital,
        'grid_pnl_usd':        round(grid_pnl, 2),
        'grid_return_pct':     round(grid_pnl / total_capital * 100, 2),
        'completed_pairs':     n_completed,
        'avg_pnl_per_pair':    round(grid_pnl / max(n_completed, 1), 2),
        'buy_and_hold_pnl':    round(bh_pnl, 2),
        'price_change_pct':    round(price_change, 2),
    }


perf = compute_grid_performance(trades_df, prices, ORDER_SIZE_USD, N_LEVELS)
print('Grid Performance Summary:')
for k, v in perf.items():
    print(f'  {k}: {v}')

def plot_grid_performance(
    trades_df: pd.DataFrame,
    prices: np.ndarray,
    grid: list,
    grid_center: float
) -> None:
    """
    Two-panel visualization: price with grid levels, and cumulative PnL.

    Parameters
    ----------
    trades_df : pd.DataFrame  Trade log.
    prices : np.ndarray  Price series.
    grid : list  List of GridLevel objects.
    grid_center : float  Grid centre price.
    """
    fig, axes = plt.subplots(2, 1, figsize=(14, 10))

    # Panel 1: Price + grid levels + trade markers
    axes[0].plot(prices, color='steelblue', linewidth=1.0, label='Price')
    for lvl in grid:
        axes[0].axhline(lvl.price, color='lightgray', linewidth=0.4, linestyle='-')
    if len(trades_df):
        buys  = trades_df[trades_df['side'] == 'BUY']
        sells = trades_df[trades_df['side'] == 'SELL']
        axes[0].scatter(buys['bar'],  buys['fill_price'],  color='green', s=15, alpha=0.6, label='Buy fills')
        axes[0].scatter(sells['bar'], sells['fill_price'], color='red',   s=15, alpha=0.6, label='Sell fills')
    axes[0].set_ylabel('Price (USD)')
    axes[0].set_title(f'{GRID_TYPE.capitalize()} Grid — {N_LEVELS} levels ±{GRID_RANGE_PCT*100:.0f}% range')
    axes[0].legend()

    # Panel 2: Cumulative PnL
    if len(trades_df):
        cum_pnl = trades_df['pnl_usd'].cumsum()
        axes[1].plot(trades_df['bar'], cum_pnl, color='green', linewidth=1.5, label='Grid PnL')
    axes[1].axhline(0, color='gray', linewidth=0.5, linestyle='--')
    axes[1].set_xlabel('Bar (day)')
    axes[1].set_ylabel('Cumulative PnL (USD)')
    axes[1].set_title('Cumulative Grid Trading PnL (Realized, from completed pairs)')
    axes[1].legend()

    plt.tight_layout()
    plt.show()


plot_grid_performance(trades_df, prices, grid, GRID_CENTER)

def export_grid_results(trades_df: pd.DataFrame, perf: dict) -> None:
    """
    Export trade log and performance summary.

    Parameters
    ----------
    trades_df : pd.DataFrame  Trade log.
    perf : dict  Performance metrics dict.
    """
    trades_df.to_csv('grid_trades.csv', index=False)
    pd.DataFrame([perf]).to_csv('grid_performance.csv', index=False)
    print('Exported: grid_trades.csv, grid_performance.csv')


export_grid_results(trades_df, perf)