import time
import numpy as np
import sys
from pathlib import Path
from unittest.mock import MagicMock

# Correct sys.path
ROOT_DIR = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT_DIR / "nautilus_dolphin"))
sys.path.insert(0, str(ROOT_DIR))

from nautilus_dolphin.nautilus.ob_features import OBFeatureEngine
from nautilus_dolphin.nautilus.ob_provider import OBSnapshot

def create_snap(asset):
    return OBSnapshot(
        timestamp=time.time(),
        asset=asset,
        bid_notional=np.random.rand(5) * 10000,
        ask_notional=np.random.rand(5) * 10000,
        bid_depth=np.random.rand(5),
        ask_depth=np.random.rand(5)
    )

def benchmark():
    provider = MagicMock()
    assets = ["BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT", "XRPUSDT"]
    provider.get_snapshot.side_effect = lambda a, t: create_snap(a)
    
    engine = OBFeatureEngine(provider)
    
    # Warmup
    for i in range(100):
        engine.step_live(assets, i)
        
    # Test
    iterations = 2000
    start = time.perf_counter()
    for i in range(iterations):
        engine.step_live(assets, 100 + i)
    end = time.perf_counter()
    
    duration = end - start
    print(f"BASELINE: {iterations} iterations in {duration:.4f}s ({iterations/duration:.2f} Hz)")
    return iterations / duration

if __name__ == "__main__":
    benchmark()