"""
Synthetic data generation for AMM backtesting.

Provides price path and swap generation using various stochastic models.
"""

import numpy as np
import pandas as pd
from typing import Optional, Literal


def generate_price_path(
    n_steps: int = 10000,
    initial_price: float = 1.0,
    volatility: float = 0.02,
    drift: float = 0.0,
    dt: float = 1.0,
    model: Literal['gbm', 'jump', 'ou'] = 'gbm',
    seed: Optional[int] = None,
    # Jump-diffusion parameters
    jump_intensity: float = 0.1,
    jump_mean: float = 0.0,
    jump_std: float = 0.05,
    # Ornstein-Uhlenbeck parameters
    mean_reversion_speed: float = 0.1,
    long_term_mean: Optional[float] = None,
) -> np.ndarray:
    """
    Generate synthetic price paths using various stochastic models.

    Parameters
    ----------
    n_steps : int
        Number of price points to generate
    initial_price : float
        Starting price
    volatility : float
        Annualized volatility (sigma) for GBM/jump, or instantaneous volatility for OU
    drift : float
        Annualized drift (mu) for GBM
    dt : float
        Time step size (in arbitrary units, affects scaling)
    model : str
        Price model to use:
        - 'gbm': Geometric Brownian Motion (default)
        - 'jump': Merton jump-diffusion
        - 'ou': Ornstein-Uhlenbeck (mean-reverting)
    seed : int, optional
        Random seed for reproducibility
    jump_intensity : float
        For 'jump' model: average number of jumps per time step (lambda)
    jump_mean : float
        For 'jump' model: mean of log-jump size
    jump_std : float
        For 'jump' model: std of log-jump size
    mean_reversion_speed : float
        For 'ou' model: speed of mean reversion (theta)
    long_term_mean : float, optional
        For 'ou' model: long-term mean price (defaults to initial_price)

    Returns
    -------
    np.ndarray
        Array of prices with shape (n_steps,)

    Examples
    --------
    >>> prices = generate_price_path(1000, volatility=0.02, model='gbm')
    >>> prices = generate_price_path(1000, model='jump', jump_intensity=0.05)
    >>> prices = generate_price_path(1000, model='ou', mean_reversion_speed=0.5)
    """
    if seed is not None:
        np.random.seed(seed)

    prices = np.zeros(n_steps)
    prices[0] = initial_price

    if model == 'gbm':
        # Geometric Brownian Motion
        # dS = mu*S*dt + sigma*S*dW
        # S(t+dt) = S(t) * exp((mu - 0.5*sigma^2)*dt + sigma*sqrt(dt)*Z)
        z = np.random.standard_normal(n_steps - 1)
        log_returns = (drift - 0.5 * volatility**2) * dt + volatility * np.sqrt(dt) * z
        prices[1:] = initial_price * np.exp(np.cumsum(log_returns))

    elif model == 'jump':
        # Merton Jump-Diffusion Model
        # dS/S = mu*dt + sigma*dW + (J-1)*dN
        # where J is log-normally distributed and N is Poisson process
        for i in range(1, n_steps):
            # Diffusion component (GBM)
            z = np.random.standard_normal()
            diffusion = (drift - 0.5 * volatility**2) * dt + volatility * np.sqrt(dt) * z

            # Jump component
            n_jumps = np.random.poisson(jump_intensity * dt)
            if n_jumps > 0:
                jump_sizes = np.random.normal(jump_mean, jump_std, n_jumps)
                jump = np.sum(jump_sizes)
            else:
                jump = 0.0

            prices[i] = prices[i - 1] * np.exp(diffusion + jump)

    elif model == 'ou':
        # Ornstein-Uhlenbeck process (mean-reverting)
        # dX = theta*(mu - X)*dt + sigma*dW
        # For prices, we apply OU to log-price to keep prices positive
        if long_term_mean is None:
            long_term_mean = initial_price

        log_mean = np.log(long_term_mean)
        log_price = np.log(initial_price)

        for i in range(1, n_steps):
            z = np.random.standard_normal()
            # Mean reversion in log space
            log_price = (
                log_price
                + mean_reversion_speed * (log_mean - log_price) * dt
                + volatility * np.sqrt(dt) * z
            )
            prices[i] = np.exp(log_price)

    else:
        raise ValueError(f"Unknown model: {model}. Use 'gbm', 'jump', or 'ou'")

    return prices


def generate_swaps(
    n_swaps: int = 10000,
    initial_price: float = 1.0,
    base_volume: float = 1000.0,
    volume_volatility: float = 0.5,
    buy_sell_ratio: float = 0.5,
    price_volatility: float = 0.02,
    price_model: Literal['gbm', 'jump', 'ou'] = 'gbm',
    start_timestamp: int = 1700000000,
    time_between_swaps: int = 15,
    seed: Optional[int] = None,
    **price_model_kwargs,
) -> pd.DataFrame:
    """
    Generate synthetic swap data for AMM backtesting.

    Parameters
    ----------
    n_swaps : int
        Number of swaps to generate
    initial_price : float
        Starting price (token1/token0)
    base_volume : float
        Base swap volume in token1 terms
    volume_volatility : float
        Volatility of swap volumes (as fraction of base_volume)
    buy_sell_ratio : float
        Probability of a buy (token0 -> token1) vs sell
    price_volatility : float
        Price volatility parameter for the underlying price model
    price_model : str
        Price model ('gbm', 'jump', 'ou')
    start_timestamp : int
        Starting Unix timestamp
    time_between_swaps : int
        Average seconds between swaps
    seed : int, optional
        Random seed for reproducibility
    **price_model_kwargs
        Additional arguments passed to generate_price_path()

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - amount0 (float64): Token0 amount (positive = into pool)
        - amount1 (float64): Token1 amount (positive = into pool)
        - price (float64): Price at swap time
        - timestamp (int64): Unix timestamp
        - volume (float64): Swap volume in token1 terms

    Examples
    --------
    >>> swaps = generate_swaps(10000, price_volatility=0.02)
    >>> swaps = generate_swaps(5000, buy_sell_ratio=0.6, price_model='jump')
    """
    if seed is not None:
        np.random.seed(seed)

    # Filter out any volatility kwarg to avoid conflicts (use price_volatility instead)
    filtered_kwargs = {k: v for k, v in price_model_kwargs.items() if k != 'volatility'}

    # Generate price path
    prices = generate_price_path(
        n_steps=n_swaps,
        initial_price=initial_price,
        volatility=price_volatility,
        model=price_model,
        seed=None,  # Already seeded above
        **filtered_kwargs,
    )

    # Generate volumes (log-normal distribution for realistic volume patterns)
    log_volumes = np.random.normal(
        np.log(base_volume),
        volume_volatility,
        n_swaps
    )
    volumes = np.exp(log_volumes)

    # Generate buy/sell directions
    is_buy = np.random.random(n_swaps) < buy_sell_ratio

    # Calculate swap amounts
    # Convention: positive amount = tokens going INTO the pool
    # Buy (token0 -> token1): amount0 > 0, amount1 < 0
    # Sell (token1 -> token0): amount0 < 0, amount1 > 0
    amount0 = np.zeros(n_swaps)
    amount1 = np.zeros(n_swaps)

    for i in range(n_swaps):
        if is_buy[i]:
            # Buying token1 with token0
            # volume is in token1 terms, so amount1_out = volume
            # amount0_in = volume / price (approximately, ignoring slippage)
            amount0[i] = volumes[i] / prices[i]  # token0 in
            amount1[i] = -volumes[i]  # token1 out (negative)
        else:
            # Selling token1 for token0
            # amount1_in = volume
            # amount0_out = volume * price (approximately)
            amount0[i] = -volumes[i] * prices[i]  # token0 out (negative)
            amount1[i] = volumes[i]  # token1 in

    # Generate timestamps
    time_deltas = np.random.exponential(time_between_swaps, n_swaps)
    timestamps = start_timestamp + np.cumsum(time_deltas).astype(np.int64)

    # Create DataFrame
    df = pd.DataFrame({
        'amount0': amount0.astype(np.float64),
        'amount1': amount1.astype(np.float64),
        'price': prices.astype(np.float64),
        'timestamp': timestamps.astype(np.int64),
        'volume': volumes.astype(np.float64),
    })

    return df
