#!/usr/bin/env python3
"""
Ultimate Raydium MEV Bot - Optimized for NVIDIA 4070 Ti
===============================================================
Features:
- CUDA 11.8 acceleration with PyTorch FP16 precision
- Real-time WebSocket balance monitoring
- Raydium V4 AMM instruction handling (0x09)
- Proper WSOL account handling
- Atomic transaction execution
- Advanced token safety validation
- Multi-threaded operation with asyncio
"""
import asyncio
import base64
import json
import logging
import os
import sys
import time
import traceback
import websockets
import torch
import numpy as np
from typing import Dict, List, Any, Optional, Set, Tuple
from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor

# Import Rust bindings - critical for high-performance execution
from solana_rust_bot import (
    RaydiumTradingCore,
    SwapParams,
    TokenPairSetting,
    TradingConfig,
    TokenSafetyCheck,
    TransactionStatus
)

#==============================================================================
# Configuration Constants - Load from environment variables
#==============================================================================

# Network constants
LAMPORTS_PER_SOL = 1_000_000_000
RAYDIUM_SWAP_INSTRUCTION = 0x09  # V4 AMM instruction code
RAYDIUM_PROGRAM_ID = os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8")
WSOL_ADDRESS = os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
TOKEN_PROGRAM_ID = os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
SYSTEM_PROGRAM_ID = os.getenv("SYSTEM_PROGRAM_ID", "11111111111111111111111111111111")

# RPC connection settings
RPC_URL = os.getenv("RPC_URL", "https://winny-rychu7-fast-mainnet.helius-rpc.com")
WS_URL = os.getenv("WS_URL", "wss://winny-rychu7-fast-mainnet.helius-rpc.com")
HELIUS_WS_URL = os.getenv("HELIUS_WS_URL", "wss://winny-rychu7-fast-mainnet.helius-rpc.com")
API_KEY = os.getenv("API_KEY", "86ab5634-df30-4a5b-bcfb-3f53d7095ca2")

# Wallet configuration
KEYPAIR_PATH = r"C:\solana_rust_bot\keypair.bin"  # Windows path with r-string
WALLET_PUBKEY = os.getenv("SOLANA_PUBLIC_KEY", "77ssrQ7LmPso1d1wYp4jAQdz4NqKSbKK8t6mFLcx6Xwr")
WSOL_TOKEN_ACCOUNT = os.getenv("WSOL_TOKEN_ACCOUNT", "5TwPK71xyhpkxNQ3WzxvSrpVpj16eDETwmMQJX74BGRX")

# Trading parameters
FRONT_RUN_AMOUNT = float(os.getenv("FRONT_RUN_AMOUNT", "0.3"))
MIN_PROFIT = float(os.getenv("MIN_PROFIT", "0.002"))
SLIPPAGE_BPS = int(os.getenv("SLIPPAGE_BPS", "1000"))
SANDWICH_MIN_VICTIM_SIZE = float(os.getenv("SANDWICH_MIN_VICTIM_SIZE", "5.0"))
SANDWICH_MAX_VICTIM_SIZE = float(os.getenv("SANDWICH_MAX_VICTIM_SIZE", "320.0"))
MIN_SOL_BALANCE = float(os.getenv("MIN_SOL_BALANCE", "0.03"))
MIN_WSOL_BALANCE = float(os.getenv("MIN_WSOL_BALANCE", "0.02"))
EXIT_DELAY = float(os.getenv("EXIT_DELAY", "3.0"))
COMPUTE_LIMIT = int(os.getenv("COMPUTE_LIMIT", "50000"))
PRIORITY_FEE = int(os.getenv("PRIORITY_FEE", "950000"))
SKIP_PREFLIGHT = os.getenv("SKIP_PREFLIGHT", "true").lower() == "true"
MAX_RETRIES = int(os.getenv("MAX_RETRIES", "2"))
MAX_TRADES_PER_MINUTE = int(os.getenv("MAX_TRADES_PER_MINUTE", "30"))

# Security settings
VERIFY_MINT_AUTHORITY = os.getenv("VERIFY_MINT_AUTHORITY", "true").lower() == "true" 
VERIFY_LIQUIDITY_LOCK = os.getenv("VERIFY_LIQUIDITY_LOCK", "true").lower() == "true"
MIN_EXIT_LIQUIDITY_SOL = float(os.getenv("MIN_EXIT_LIQUIDITY_SOL", "10.0"))
SECURITY_CACHE_SECONDS = int(os.getenv("SECURITY_CACHE_SECONDS", "600"))

# GPU settings
GPU_MEMORY_FRACTION = float(os.getenv("GPU_MEMORY_FRACTION", "0.95"))
TENSOR_SIZE = int(os.getenv("TENSOR_SIZE", "176"))
STREAM_PRIORITY = int(os.getenv("STREAM_PRIORITY", "-3"))
USE_AMP = os.getenv("USE_AMP", "true").lower() == "true"
ALLOW_TF32 = os.getenv("ALLOW_TF32", "true").lower() == "true"
CUDA_BENCHMARK = os.getenv("CUDA_BENCHMARK", "true").lower() == "true"
CUDA_DETERMINISTIC = os.getenv("CUDA_DETERMINISTIC", "false").lower() == "true"
BATCH_SIZE = int(os.getenv("BATCH_SIZE", "512"))

#==============================================================================
# Logging Configuration
#==============================================================================

# Configure logging with timestamps, log level, and both console and file output
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
log_file = f"raydium_mev_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}.log"

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL),
    format='%(asctime)s | %(levelname)s | %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(log_file)
    ]
)
logger = logging.getLogger("raydium_mev")

#==============================================================================
# GPU Acceleration
#==============================================================================

class GPUAccelerator:
    """High-performance GPU acceleration optimized for 4070 Ti"""
    
    def __init__(self):
        """Initialize GPU with optimized settings"""
        # Configure device and parameters
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tensor_size = TENSOR_SIZE
        self.batch_size = BATCH_SIZE
        
        # Configure CUDA settings for maximum performance
        if torch.cuda.is_available():
            # Enable TF32 for faster computation on Ampere GPUs (if allowed)
            if ALLOW_TF32:
                torch.backends.cuda.matmul.allow_tf32 = True
                torch.backends.cudnn.allow_tf32 = True
            
            # Set cuDNN benchmark and determinism
            torch.backends.cudnn.benchmark = CUDA_BENCHMARK
            torch.backends.cudnn.deterministic = CUDA_DETERMINISTIC
            
            # Create CUDA stream with high priority
            self.stream = torch.cuda.Stream(priority=STREAM_PRIORITY)
            
            # Configure precision (FP16/FP32)
            self.dtype = torch.float16 if USE_AMP else torch.float32
            
            # Set memory fraction
            if hasattr(torch.cuda, 'set_per_process_memory_fraction'):
                torch.cuda.set_per_process_memory_fraction(GPU_MEMORY_FRACTION)
            
            # Log GPU configuration
            logger.info(f"GPU acceleration enabled: {torch.cuda.get_device_name(0)}")
            logger.info(f"CUDA version: {torch.version.cuda}")
            logger.info(f"Available memory: {torch.cuda.get_device_properties(0).total_memory / 1024 / 1024 / 1024:.2f} GB")
            logger.info(f"Using tensor size: {self.tensor_size}, batch size: {self.batch_size}")
            logger.info(f"FP16/AMP enabled: {USE_AMP}, TF32 enabled: {ALLOW_TF32}")
        else:
            self.dtype = torch.float32
            logger.warning("GPU not available, using CPU for calculations")
            
        # Initialize tensors and warm up GPU
        self.initialize_tensors()
    
    def initialize_tensors(self):
        """Initialize tensors and warm up GPU"""
        # Initialize trading calculation tensors
        self.prices = torch.zeros(self.tensor_size, dtype=self.dtype, device=self.device)
        self.amounts = torch.zeros(self.tensor_size, dtype=self.dtype, device=self.device)
        self.batch_data = torch.zeros((self.batch_size, self.tensor_size), dtype=self.dtype, device=self.device)
        
        # Initialize AMM calculation tensors
        self.pool_reserves = torch.zeros(2, dtype=self.dtype, device=self.device)
        self.amount_tensor = torch.zeros(1, dtype=self.dtype, device=self.device)
        
        # Warm up GPU with some basic operations
        if torch.cuda.is_available():
            with torch.cuda.stream(self.stream):
                for _ in range(100):
                    # Matrix operations
                    _ = torch.matmul(self.prices, self.amounts)
                    _ = torch.relu(self.batch_data)
                    
                    # AMM simulation
                    k = self.pool_reserves[0] * self.pool_reserves[1]
                    new_reserves = self.pool_reserves + self.amount_tensor
                    _ = k / new_reserves
                
                # Ensure all operations complete
                torch.cuda.synchronize()
                logger.debug("GPU warm-up complete")
    
    @torch.compile(fullgraph=True)
    def calculate_price_impact(self, pool_reserves: Tuple[float, float], amount_in: float, 
                             is_buy: bool) -> float:
        """Calculate price impact using GPU-optimized calculations"""
        # Transfer data to existing tensors (avoids allocations)
        self.pool_reserves[0] = pool_reserves[0]
        self.pool_reserves[1] = pool_reserves[1]
        self.amount_tensor[0] = amount_in
        
        # Calculate constant product formula (x * y = k)
        k = self.pool_reserves[0] * self.pool_reserves[1]
        
        if is_buy:
            # Buying token with SOL
            new_reserve_in = self.pool_reserves[0] + self.amount_tensor[0]
            new_reserve_out = k / new_reserve_in
            amount_out = self.pool_reserves[1] - new_reserve_out
            price_impact = 1.0 - ((amount_out / self.amount_tensor[0]) * 
                                  (self.pool_reserves[0] / self.pool_reserves[1]))
        else:
            # Selling token for SOL
            new_reserve_in = self.pool_reserves[1] + self.amount_tensor[0]
            new_reserve_out = k / new_reserve_in
            amount_out = self.pool_reserves[0] - new_reserve_out
            price_impact = 1.0 - ((amount_out / self.amount_tensor[0]) * 
                                  (self.pool_reserves[1] / self.pool_reserves[0]))
        
        # Return scalar value
        return price_impact.item()
    
    def estimate_sandwich_profit(self, pool_reserves: Tuple[float, float], front_run_amount: float, 
                               victim_amount: float) -> Dict[str, float]:
        """Estimate sandwich attack profit using GPU calculations"""
        # Run calculations on GPU with proper error handling
        if torch.cuda.is_available():
            try:
                with torch.cuda.stream(self.stream):
                    with torch.cuda.amp.autocast(enabled=USE_AMP):
                        # Step 1: Our front-run buy
                        impact1 = self.calculate_price_impact(pool_reserves, front_run_amount, True)
                        buy_amount_out = front_run_amount * (1 - impact1)
                        
                        # New pool state after our buy
                        new_pool_token = pool_reserves[0] + front_run_amount
                        new_pool_sol = pool_reserves[1] - buy_amount_out
                        
                        # Step 2: Victim's trade at worse price
                        impact2 = self.calculate_price_impact((new_pool_token, new_pool_sol), 
                                                           victim_amount, True)
                        victim_amount_out = victim_amount * (1 - impact2)
                        
                        # Step 3: Our sell after victim's trade
                        final_pool_token = new_pool_token + victim_amount
                        final_pool_sol = new_pool_sol - victim_amount_out
                        
                        # Sell our tokens
                        impact3 = self.calculate_price_impact((final_pool_token, final_pool_sol), 
                                                           buy_amount_out, False)
                        sell_amount_out = buy_amount_out * (1 - impact3)
                        
                        # Calculate profit and metrics
                        profit = sell_amount_out - front_run_amount
                        roi = (profit / front_run_amount) * 100 if front_run_amount > 0 else 0
                        
                        # Synchronize to ensure calculations complete
                        torch.cuda.synchronize()
                        
                        return {
                            "profit": float(profit),
                            "roi": float(roi),
                            "buy_tokens": float(buy_amount_out),
                            "sell_sol": float(sell_amount_out),
                            "victim_impact": float(impact2 * 100),  # as percentage
                            "front_run_impact": float(impact1 * 100),  # as percentage
                            "exit_impact": float(impact3 * 100),  # as percentage
                        }
            except Exception as e:
                logger.error(f"GPU calculation error: {e}")
                logger.debug(traceback.format_exc())
                
        # Fallback to simple CPU calculation
        estimated_profit = victim_amount * 0.005  # Estimated 0.5% profit
        return {
            "profit": estimated_profit,
            "roi": 0.5,
            "buy_tokens": front_run_amount * 0.95,
            "sell_sol": front_run_amount * 1.005,
            "victim_impact": 0.5,
            "front_run_impact": 0.5,
            "exit_impact": 0.5
        }
    
    def cleanup(self):
        """Release GPU resources"""
        if torch.cuda.is_available():
            # Empty CUDA cache
            torch.cuda.empty_cache()
            
            # Synchronize all streams
            if hasattr(self, 'stream'):
                self.stream.synchronize()
                
            logger.info("GPU resources released successfully")

#==============================================================================
# WebSocket Balance Monitor
#==============================================================================

class BalanceMonitor:
    """Real-time balance monitor using WebSocket for SOL and WSOL"""
    
    def __init__(self, wallet_pubkey: str, wsol_account: str):
        self.wallet_pubkey = wallet_pubkey
        self.wsol_account = wsol_account
        self.sol_balance = 0.0
        self.wsol_balance = 0.0
        self.last_sol_balance = 0.0
        self.last_wsol_balance = 0.0
        self.is_running = False
        self.task = None
        self.callbacks = []
        
        # Track trades for profit calculation
        self.trade_history = []
        self.last_update_time = time.time()
        
    def register_callback(self, callback):
        """Register a callback for balance updates"""
        self.callbacks.append(callback)
    
    async def start(self, initial_sol: float, initial_wsol: float):
        """Start monitoring with initial balances"""
        self.sol_balance = initial_sol
        self.wsol_balance = initial_wsol
        self.last_sol_balance = initial_sol
        self.last_wsol_balance = initial_wsol
        self.is_running = True
        
        # Start WebSocket subscription in background
        self.task = asyncio.create_task(self.monitor_balances())
        logger.info(f"Balance monitor started with SOL: {initial_sol:.4f}, WSOL: {initial_wsol:.4f}")
        
        return True
    
    async def monitor_balances(self):
        """Monitor SOL and WSOL balances via WebSocket"""
        reconnect_attempts = 0
        max_reconnect_delay = 60  # Maximum backoff in seconds
        
        while self.is_running:
            try:
                # Connect to WebSocket with API key in headers
                async with websockets.connect(
                    HELIUS_WS_URL,
                    extra_headers={"api-key": API_KEY},
                    ping_interval=20,
                    ping_timeout=30,
                    close_timeout=60
                ) as ws:
                    logger.info("Connected to balance WebSocket")
                    reconnect_attempts = 0
                    
                    # Subscribe to native SOL account
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "accountSubscribe",
                        "params": [
                            self.wallet_pubkey,
                            {"encoding": "jsonParsed", "commitment": "confirmed"}
                        ]
                    }))
                    
                    # Subscribe to WSOL token account
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 2,
                        "method": "accountSubscribe",
                        "params": [
                            self.wsol_account,
                            {"encoding": "jsonParsed", "commitment": "confirmed"}
                        ]
                    }))
                    
                    # Track subscriptions 
                    subscriptions = {}
                    
                    # Get and process subscription IDs
                    sol_response = await ws.recv()
                    sol_data = json.loads(sol_response)
                    if "result" in sol_data:
                        subscriptions[sol_data["result"]] = "SOL"
                        logger.debug(f"SOL subscription ID: {sol_data['result']}")
                    
                    wsol_response = await ws.recv()
                    wsol_data = json.loads(wsol_response)
                    if "result" in wsol_data:
                        subscriptions[wsol_data["result"]] = "WSOL"
                        logger.debug(f"WSOL subscription ID: {wsol_data['result']}")
                    
                    # Process balance updates
                    while self.is_running:
                        try:
                            # Wait for message with timeout
                            msg = await asyncio.wait_for(ws.recv(), timeout=30)
                            data = json.loads(msg)
                            
                            # Handle balance updates
                            if "method" in data and data["method"] == "accountNotification":
                                sub_id = data["params"]["subscription"]
                                acc_type = subscriptions.get(sub_id)
                                
                                if acc_type == "SOL":
                                    # Process SOL balance update
                                    if "lamports" in data["params"]["result"]["value"]:
                                        lamports = data["params"]["result"]["value"]["lamports"]
                                        new_sol_balance = lamports / LAMPORTS_PER_SOL
                                        
                                        if new_sol_balance != self.sol_balance:
                                            old_balance = self.sol_balance
                                            self.last_sol_balance = self.sol_balance
                                            self.sol_balance = new_sol_balance
                                            logger.info(f"SOL balance changed: {old_balance:.4f} → {new_sol_balance:.4f}")
                                            
                                            # Record balance change
                                            self.trade_history.append({
                                                'timestamp': time.time(),
                                                'type': 'SOL',
                                                'old_balance': old_balance,
                                                'new_balance': new_sol_balance
                                            })
                                            
                                            # Notify callbacks
                                            for callback in self.callbacks:
                                                await callback('SOL', old_balance, new_sol_balance)
                                
                                elif acc_type == "WSOL":
                                    # Process WSOL balance update
                                    value = data["params"]["result"]["value"]
                                    if "data" in value and "parsed" in value["data"]:
                                        parsed_data = value["data"]["parsed"]
                                        if "info" in parsed_data and "tokenAmount" in parsed_data["info"]:
                                            token_amount = parsed_data["info"]["tokenAmount"]
                                            if "uiAmount" in token_amount:
                                                new_wsol_balance = token_amount["uiAmount"]
                                                
                                                if new_wsol_balance != self.wsol_balance:
                                                    old_balance = self.wsol_balance
                                                    self.last_wsol_balance = self.wsol_balance
                                                    self.wsol_balance = new_wsol_balance
                                                    logger.info(f"WSOL balance changed: {old_balance:.4f} → {new_wsol_balance:.4f}")
                                                    
                                                    # Record balance change
                                                    self.trade_history.append({
                                                        'timestamp': time.time(),
                                                        'type': 'WSOL',
                                                        'old_balance': old_balance,
                                                        'new_balance': new_wsol_balance
                                                    })
                                                    
                                                    # Notify callbacks
                                                    for callback in self.callbacks:
                                                        await callback('WSOL', old_balance, new_wsol_balance)
                        
                        except asyncio.TimeoutError:
                            # No update received, send ping
                            await ws.ping()
                            continue
                        
            except (websockets.exceptions.ConnectionClosed, 
                    websockets.exceptions.WebSocketException) as e:
                if not self.is_running:
                    break
                    
                # Implement backoff for reconnection attempts
                reconnect_attempts += 1
                backoff = min(5 * reconnect_attempts, max_reconnect_delay)
                logger.error(f"Balance WebSocket error: {e}. "
                           f"Reconnecting in {backoff} seconds...")
                await asyncio.sleep(backoff)
                
            except Exception as e:
                if not self.is_running:
                    break
                    
                logger.error(f"Unexpected error in balance monitor: {e}")
                logger.error(traceback.format_exc())
                await asyncio.sleep(5)
    
    def calculate_profit(self, window_seconds: float = 60.0) -> float:
        """Calculate profit over a time window based on SOL balance changes"""
        now = time.time()
        start_time = now - window_seconds
        
        # Filter trade history to the relevant time window
        relevant_trades = [t for t in self.trade_history if t['timestamp'] >= start_time]
        
        if not relevant_trades:
            return 0.0
        
        # Get the earliest and latest balances in the window
        sol_changes = [t for t in relevant_trades if t['type'] == 'SOL']
        wsol_changes = [t for t in relevant_trades if t['type'] == 'WSOL']
        
        if not sol_changes:
            return 0.0
            
        # Sort by timestamp
        sol_changes.sort(key=lambda x: x['timestamp'])
        
        # Calculate net change over the period
        net_sol_change = sol_changes[-1]['new_balance'] - sol_changes[0]['old_balance']
        
        # Return estimated profit (fees already accounted for in balance)
        return net_sol_change
    
    async def stop(self):
        """Stop the balance monitor"""
        if not self.is_running:
            return
            
        self.is_running = False
        
        if self.task:
            try:
                self.task.cancel()
                await asyncio.wait_for(asyncio.shield(self.task), timeout=5)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pass
            
        logger.info("Balance monitor stopped")

#==============================================================================
# Main MEV Bot Implementation
#==============================================================================

class RaydiumMEVBot:
    """Ultimate MEV bot for front-running Raydium swaps with GPU acceleration"""
    
    def __init__(self):
        """Initialize MEV bot with all components"""
        logger.info("Initializing Raydium MEV Bot")
        
        # Initialize trading core with correct keypair path
        self.trading_core = RaydiumTradingCore(
            rpc_url=RPC_URL,
            keypair_path=KEYPAIR_PATH
        )
        
        # Get wallet public key
        self.wallet_pubkey = self.trading_core.get_public_key()
        logger.info(f"Wallet public key: {self.wallet_pubkey}")
        
        # Initialize GPU accelerator
        self.gpu = GPUAccelerator()
        
        # State tracking
        self.active_trades = set()  # Currently executing tx signatures
        self.seen_txs = set()       # Already processed tx signatures
        self.token_safety_cache = {}  # Cache of token safety checks
        self.current_position = None  # Current open position
        self.token_accounts = {}    # Map of token mint to ATA
        self.last_trade_time = 0    # Last trade timestamp
        self.trades_this_minute = 0 # Rate limiting
        self.minute_start = time.time()
        
        # Statistics tracking
        self.stats = {
            'success_count': 0,
            'failure_count': 0,
            'profit_total': 0.0,
            'last_profit': 0.0,
            'trades_today': 0,
            'start_time': time.time(),
            'last_sol_balance': 0.0
        }
        
        # Balance monitor
        self.balance_monitor = None
        
        # Initialize with healthcheck
        self.is_running = True
        self.shutdown_event = asyncio.Event()
        
        # Thread pool for CPU-bound operations
        self.thread_pool = ThreadPoolExecutor(max_workers=4)
        
    async def initialize(self):
        """Initialize the bot with balance checks and account setup"""
        try:
            logger.info("Performing initial setup...")
            
            # Check balances
            balances = await self.trading_core.get_balances()
            sol_balance, token_balances = balances
            
            logger.info(f"Initial SOL Balance: {sol_balance:.4f}")
            self.stats['last_sol_balance'] = sol_balance
            
            # Initialize token accounts dict
            self.token_accounts = {}
            for token_mint, amount in token_balances.items():
                if amount > 0:
                    # Ensure ATA exists and get the address
                    ata = await self.trading_core.ensure_ata(token_mint)
                    self.token_accounts[token_mint] = ata
                    logger.info(f"Token {token_mint}: {amount:.4f} (Account: {ata})")
            
            # Ensure WSOL ATA exists
            wsol_ata = await self.trading_core.ensure_ata(WSOL_ADDRESS)
            self.token_accounts[WSOL_ADDRESS] = wsol_ata
            wsol_balance = token_balances.get(WSOL_ADDRESS, 0)
            logger.info(f"WSOL Account: {wsol_ata}, Balance: {wsol_balance:.4f}")
            
            # Initialize balance monitor with the accounts
            self.balance_monitor = BalanceMonitor(
                wallet_pubkey=self.wallet_pubkey,
                wsol_account=wsol_ata if wsol_ata else WSOL_TOKEN_ACCOUNT
            )
            
            # Register balance change callback
            self.balance_monitor.register_callback(self.handle_balance_change)
            
            # Start balance monitor
            await self.balance_monitor.start(sol_balance, wsol_balance)
            
            # Start healthcheck in background
            asyncio.create_task(self.healthcheck())
            
            logger.info("Bot initialization complete and ready for trading")
            return True
            
        except Exception as e:
            logger.error(f"Initialization error: {e}")
            logger.error(traceback.format_exc())
            return False
        
    async def handle_balance_change(self, balance_type: str, old_value: float, new_value: float):
        """Handle balance change notification from WebSocket"""
        try:
            # If we have an open position and token balance dropped, it indicates a successful exit
            if self.current_position and balance_type == 'WSOL' and old_value > new_value:
                # This might be our exit transaction completing
                token_mint = self.current_position.get('token_mint')
                
                if token_mint and token_mint == WSOL_ADDRESS:
                    logger.info(f"Detected WSOL position exit: {old_value:.4f} → {new_value:.4f}")
                    
                    # Verify if SOL balance increased correspondingly
                    sol_change = self.balance_monitor.sol_balance - self.balance_monitor.last_sol_balance
                    
                    if sol_change > 0:
                        # Calculate profit based on SOL balance increase
                        profit = sol_change
                        logger.info(f"Estimated profit: {profit:.6f} SOL")
                        
                        # Update stats
                        self.stats['last_profit'] = profit
                        self.stats['profit_total'] += profit
                        self.stats['success_count'] += 1
                        
                        # Clear current position
                        self.current_position = None
                        
            # General SOL balance increase might indicate profit from other trades
            if balance_type == 'SOL' and new_value > old_value:
                # Only record significant balance changes
                change = new_value - old_value
                if change > 0.001 and not self.current_position:
                    logger.info(f"Detected SOL balance increase: +{change:.6f} SOL")
                    
                    # Record this as potential profit if not associated with a specific trade
                    if time.time() - self.last_trade_time > 10:
                        self.stats['profit_total'] += change
                        logger.info(f"Added {change:.6f} SOL to total profit")
                        
        except Exception as e:
            logger.error(f"Error handling balance change: {e}")
            logger.debug(traceback.format_exc())

    async def healthcheck(self):
        """Periodic healthcheck to monitor system status"""
        try:
            while self.is_running:
                # Calculate uptime
                uptime = time.time() - self.stats['start_time']
                hours = uptime / 3600
                
                # Get current balances
                sol_balance = self.balance_monitor.sol_balance if self.balance_monitor else 0
                wsol_balance = self.balance_monitor.wsol_balance if self.balance_monitor else 0
                
                logger.info(
                    f"Bot status: Running for {hours:.2f}h | "
                    f"SOL: {sol_balance:.4f} | WSOL: {wsol_balance:.4f} | "
                    f"Active trades: {len(self.active_trades)} | "
                    f"Position: {'Open' if self.current_position else 'None'} | "
                    f"Success/Fail: {self.stats['success_count']}/{self.stats['failure_count']} | "
                    f"Profit: {self.stats['profit_total']:.6f} SOL"
                )
                
                # Check GPU health if using CUDA
                if torch.cuda.is_available():
                    mem_used = torch.cuda.memory_allocated() / 1024 / 1024 / 1024
                    mem_reserved = torch.cuda.memory_reserved() / 1024 / 1024 / 1024
                    logger.debug(f"GPU memory: {mem_used:.2f} GB used, {mem_reserved:.2f} GB reserved")
                
                # Reset rate limiting if needed
                now = time.time()
                if now - self.minute_start > 60:
                    self.trades_this_minute = 0
                    self.minute_start = now
                
                await asyncio.sleep(30)
                
        except asyncio.CancelledError:
            logger.info("Healthcheck task cancelled")
        except Exception as e:
            logger.error(f"Healthcheck error: {e}")
            logger.error(traceback.format_exc())

    async def monitor_mempool(self):
        """Monitor Solana mempool for Raydium swap transactions"""
        logger.info(f"Starting mempool monitoring with connection to {HELIUS_WS_URL}")
        
        # Use Helius WebSocket for optimized performance
        ws_url = HELIUS_WS_URL
        
        # Configure headers with API key
        headers = {"api-key": API_KEY} if API_KEY else {}
        headers["User-Agent"] = "solana-client/1.18.23 (client:SolanaLabs)"
        
        reconnect_attempts = 0
        max_reconnect_delay = 60  # Maximum backoff in seconds
        
        while self.is_running:
            try:
                # Connect to WebSocket
                async with websockets.connect(ws_url, extra_headers=headers) as ws:
                    logger.info("Connected to mempool WebSocket")
                    reconnect_attempts = 0
                    
                    # Subscribe to Raydium program logs
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "logsSubscribe",
                        "params": [
                            {"mentions": [RAYDIUM_PROGRAM_ID]},
                            {"commitment": "processed"}  # Fastest commitment level
                        ]
                    }))
                    
                    # Wait for subscription confirmation
                    response = await asyncio.wait_for(ws.recv(), timeout=10)
                    data = json.loads(response)
                    
                    if "result" in data:
                        sub_id = data["result"]
                        logger.info(f"Mempool subscription established: ID {sub_id}")
                    else:
                        logger.error(f"Failed to establish mempool subscription: {data}")
                        if "error" in data:
                            logger.error(f"Error: {data['error']}")
                        continue
                    
                    # Process incoming messages
                    while self.is_running:
                        try:
                            # Wait for message with timeout
                            msg = await asyncio.wait_for(ws.recv(), timeout=30)
                            await self.process_transaction(msg)
                        except asyncio.TimeoutError:
                            # No message received, send ping to keep connection alive
                            await ws.ping()
                            continue
            
            except (websockets.exceptions.ConnectionClosed, 
                    websockets.exceptions.WebSocketException) as e:
                if not self.is_running:
                    break
                    
                # Implement backoff for reconnection attempts
                reconnect_attempts += 1
                backoff = min(5 * reconnect_attempts, max_reconnect_delay)
                logger.error(f"Mempool WebSocket error: {e}. "
                           f"Reconnecting in {backoff} seconds...")
                await asyncio.sleep(backoff)
                
            except Exception as e:
                if not self.is_running:
                    break
                    
                logger.error(f"Unexpected error in mempool monitor: {e}")
                logger.error(traceback.format_exc())
                await asyncio.sleep(5)

    async def process_transaction(self, msg: str):
        """Process incoming transaction data with GPU-accelerated analysis"""
        try:
            data = json.loads(msg)
            
            # Check if this is a proper logs notification
            if "method" not in data or data["method"] != "logsNotification":
                return
                
            # Extract logs and signature
            logs = data.get('params', {}).get('result', {}).get('value', {}).get('logs', [])
            signature = data.get('params', {}).get('result', {}).get('value', {}).get('signature', '')
            
            # Skip if no logs or already processed this transaction
            if not logs or signature in self.seen_txs:
                return
            
            # Add to seen transactions
            self.seen_txs.add(signature)
            
            # Check if this is our own transaction
            if signature in self.active_trades:
                logger.info(f"Received confirmation for our transaction: {signature}")
                return
                
            # Skip if we already have a position
            if self.current_position:
                return
                
            # Enforce rate limiting
            if self.trades_this_minute >= MAX_TRADES_PER_MINUTE:
                return
                
            # Quickly check for Raydium program and swap instruction
            has_raydium_program = False
            has_swap_instruction = False
            
            for log in logs:
                if f"Program {RAYDIUM_PROGRAM_ID} invoke" in log:
                    has_raydium_program = True
                if "Program log: Instruction: Swap" in log:
                    has_swap_instruction = True
                    break
            
            # If potential swap found, extract details and check profitability
            if has_raydium_program and has_swap_instruction:
                # Extract swap details with optimized parsing
                swap_data = await self.parse_swap_instruction(logs, signature)
                
                if swap_data:
                    # Check profitability asynchronously with a short timeout
                    is_profitable_task = asyncio.create_task(self.is_profitable(swap_data))
                    try:
                        profitable = await asyncio.wait_for(is_profitable_task, timeout=0.05)
                        if profitable:
                            # Execute front-run if profitable
                            await self.execute_front_run(swap_data, signature)
                    except asyncio.TimeoutError:
                        # If profitability check takes too long, move on
                        is_profitable_task.cancel()

        except json.JSONDecodeError:
            # Not valid JSON, ignore silently
            pass
        except Exception as e:
            logger.error(f"Transaction processing error: {e}")
            logger.debug(traceback.format_exc())

    async def parse_swap_instruction(self, logs: List[str], signature: str) -> Optional[Dict]:
        """Extract swap details from transaction logs with optimized parsing"""
        try:
            # Fast path: extract key details with minimal processing
            token_mint = None
            amount_in_sol = None
            pool_id = None
            
            # Track instruction details
            is_raydium_v4 = False
            has_swap_instruction = False
            has_instruction_9 = False
            
            # Extract instruction data (loop through logs once)
            for log in logs:
                # Check for Raydium program invoke (V4)
                if f"Program {RAYDIUM_PROGRAM_ID} invoke" in log:
                    is_raydium_v4 = True
                    continue
                    
                # Check for explicit swap instruction
                if "Program log: Instruction: Swap" in log and is_raydium_v4:
                    has_swap_instruction = True
                    continue
                    
                # Check for AMM ID (pool identifier)
                if is_raydium_v4 and "Pool" in log and len(log.strip()) >= 32:
                    potential_pool = log.strip()
                    if all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in potential_pool):
                        pool_id = potential_pool
                    continue
                    
                # Extract amount from instruction data
                if "Program data: " in log and is_raydium_v4:
                    try:
                        data = base64.b64decode(log.split("Program data: ")[1])
                        if len(data) > 0 and data[0] == RAYDIUM_SWAP_INSTRUCTION:  # Check for 0x09
                            has_instruction_9 = True
                            
                            # Extract amount from first 8 bytes after instruction type
                            if len(data) >= 9:
                                amount_lamports = int.from_bytes(data[1:9], byteorder='little')
                                amount_in_sol = amount_lamports / LAMPORTS_PER_SOL
                                logger.debug(f"Extracted amount: {amount_in_sol:.4f} SOL")
                    except Exception as e:
                        logger.debug(f"Error extracting instruction data: {e}")
                    continue
                
                # Look for token mint in token program logs
                if token_mint is None and TOKEN_PROGRAM_ID in log:
                    parts = log.split()
                    for i, part in enumerate(parts):
                        # Valid token mint addresses are base58 encoded and 32+ chars
                        if len(part) >= 32 and all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in part):
                            # Skip if it's our wallet or WSOL (might be false positive)
                            if part != self.wallet_pubkey and part != WSOL_ADDRESS:
                                token_mint = part
                                break
            
            # Skip if not a valid V4 AMM transaction with instruction 9
            if not is_raydium_v4 or not has_instruction_9:
                return None
                
            # If we couldn't find a token, try harder (sometimes embedded in other logs)
            if token_mint is None:
                for log in logs:
                    if "token:" in log.lower() or "Token:" in log:
                        parts = log.split()
                        for part in parts:
                            if len(part) >= 32 and all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in part):
                                if part != self.wallet_pubkey and part != WSOL_ADDRESS:
                                    token_mint = part
                                    break
                    if token_mint:
                        break
            
            # Skip if we couldn't identify a token
            if not token_mint:
                return None
                
            # Normalize SOL address to WSOL if needed (handle system program address)
            if token_mint == SYSTEM_PROGRAM_ID:
                token_mint = WSOL_ADDRESS
                
            # Use default amount if we couldn't extract it
            if not amount_in_sol:
                amount_in_sol = 1.0  # Conservative default
                
            # Create swap data structure with all extracted information
            swap_data = {
                "token_mint": token_mint,
                "amount_in_sol": amount_in_sol,
                "signature": signature,
                "direction": "buy",  # Assume buying tokens with SOL
                "timestamp": time.time()
            }
            
            # Add pool ID if found (useful for multi-pool tokens)
            if pool_id:
                swap_data["pool_id"] = pool_id
                
            logger.debug(f"Found potential swap: {token_mint}, {amount_in_sol:.4f} SOL")
            return swap_data
            
        except Exception as e:
            logger.error(f"Error parsing swap instruction: {e}")
            logger.debug(traceback.format_exc())
            return None
            
    async def is_profitable(self, swap_data: Dict) -> bool:
        """Determine if a front-running opportunity is profitable"""
        try:
            # Extract key data
            token_mint = swap_data.get("token_mint")
            victim_amount = swap_data.get("amount_in_sol", 1.0)
            
            # Skip transactions with missing data
            if not token_mint:
                return False
                
            # Check if transaction is too recent (prevent double-front-running)
            now = time.time()
            if now - self.last_trade_time < 0.5:  # 500ms cooldown
                logger.debug("Transaction too soon after last trade")
                return False
                
            # Check token safety
            # Use cached result if available
            current_time = time.time()
            if token_mint in self.token_safety_cache:
                cache_time, safety = self.token_safety_cache[token_mint]
                if current_time - cache_time < SECURITY_CACHE_SECONDS:
                    if not safety.is_safe:
                        logger.debug(f"Token {token_mint} failed cached safety check")
                        return False
                else:              # Cache expired, remove it
                    del self.token_safety_cache[token_mint]
            
            # Not in cache or expired, perform safety check
            if token_mint not in self.token_safety_cache:
                safety = await self.trading_core.check_token_safety(token_mint)
                self.token_safety_cache[token_mint] = (current_time, safety)
                
                if not safety.is_safe:
                    logger.warning(f"Token {token_mint} failed safety check: {safety.reason}")
                    return False
                    
                # Check if the token has sufficient liquidity
                if safety.liquidity_sol < MIN_EXIT_LIQUIDITY_SOL:
                    logger.warning(f"Insufficient liquidity for {token_mint}: {safety.liquidity_sol} SOL")
                    return False
                
                # Check mint authority revoked if required
                if VERIFY_MINT_AUTHORITY and not safety.auth_revoked:
                    logger.warning(f"Mint authority not revoked for {token_mint}")
                    return False
            
            # Check if victim's trade is large enough to be worth front-running
            if victim_amount < SANDWICH_MIN_VICTIM_SIZE:
                logger.debug(f"Victim amount too small: {victim_amount} SOL < {SANDWICH_MIN_VICTIM_SIZE} SOL")
                return False
                
            if victim_amount > SANDWICH_MAX_VICTIM_SIZE:
                logger.debug(f"Victim amount too large: {victim_amount} SOL > {SANDWICH_MAX_VICTIM_SIZE} SOL")
                return False
            
            # Get current balances to ensure we have enough SOL
            if self.balance_monitor:
                sol_balance = self.balance_monitor.sol_balance
            else:
                # Fallback to direct balance check
                balances = await self.trading_core.get_balances()
                sol_balance, _ = balances
            
            # Check if we have enough SOL for front-running
            if sol_balance < FRONT_RUN_AMOUNT + MIN_SOL_BALANCE:
                logger.warning(f"Insufficient SOL balance: {sol_balance} SOL < {FRONT_RUN_AMOUNT + MIN_SOL_BALANCE} SOL")
                return False
            
            # Use GPU for profit estimation if available
            estimated_profit = 0.0
            if hasattr(self, 'gpu'):
                # Assume conservative pool reserves based on victim size
                # In production, you would fetch actual pool reserves from the chain
                pool_reserves = (victim_amount * 20, victim_amount * 20)  # Symmetric pool assumption
                
                # Calculate potential profit with GPU acceleration
                result = self.gpu.estimate_sandwich_profit(
                    pool_reserves=pool_reserves,
                    front_run_amount=FRONT_RUN_AMOUNT,
                    victim_amount=victim_amount
                )
                
                estimated_profit = result.get("profit", 0.0)
                roi_percentage = result.get("roi", 0.0)
                
                # Log detailed metrics for high-value opportunities
                if estimated_profit > MIN_PROFIT:
                    logger.info(f"Profitable opportunity: {token_mint}")
                    logger.info(f"Estimated profit: {estimated_profit:.6f} SOL ({roi_percentage:.2f}% ROI)")
                    logger.info(f"Victim size: {victim_amount:.4f} SOL")
            else:
                # Fallback profit estimation without GPU
                if victim_amount < 10.0:
                    estimated_profit = victim_amount * 0.005  # 0.5%
                elif victim_amount < 50.0:
                    estimated_profit = victim_amount * 0.004  # 0.4%
                else:
                    estimated_profit = victim_amount * 0.003  # 0.3%
            
            # Check minimum profit threshold
            if estimated_profit < MIN_PROFIT:
                logger.debug(f"Estimated profit too low: {estimated_profit:.6f} SOL < {MIN_PROFIT} SOL")
                return False
                
            # If we've passed all checks, this opportunity is profitable
            return True
            
        except Exception as e:
            logger.error(f"Error checking profitability: {e}")
            logger.debug(traceback.format_exc())
            return False

    async def execute_front_run(self, swap_data: Dict, victim_signature: str):
        """Execute front-running strategy with optimized transaction handling"""
        # Skip if already in a trade or already processed this victim
        if self.current_position or victim_signature in self.active_trades:
            return

        token_mint = swap_data.get("token_mint")
        if not token_mint:
            logger.error("Missing token mint in swap data")
            return
            
        # Rate limiting check
        now = time.time()
        if now - self.minute_start > 60:
            self.trades_this_minute = 0
            self.minute_start = now
            
        if self.trades_this_minute >= MAX_TRADES_PER_MINUTE:
            logger.warning(f"Rate limit exceeded: {self.trades_this_minute}/{MAX_TRADES_PER_MINUTE} trades this minute")
            return
            
        # Add to active trades to prevent duplicates
        self.active_trades.add(victim_signature)
        self.last_trade_time = now
        self.trades_this_minute += 1
        
        try:
            logger.info(f"Executing front-run for {token_mint}")
            logger.info(f"Victim transaction: {victim_signature}")
            
            # Generate correct instruction data for Raydium V4 AMM
            instruction_data = await asyncio.to_thread(
                self.trading_core.create_amm_instruction_data,
                FRONT_RUN_AMOUNT,  # amount_sol
                None  # min_amount_out (None = 0)
            )
            
            # Ensure WSOL ATA exists and is funded
            wsol_ata = await self.trading_core.ensure_ata(WSOL_ADDRESS)
            logger.debug(f"Using WSOL ATA: {wsol_ata}")
            
            # Record start time for performance tracking
            start_time = time.time()
            
            # Execute the front-running swap with high priority
            tx_status = await self.trading_core.execute_custom_swap(
                token_mint=token_mint,
                instruction_data=instruction_data
            )

            # Calculate execution time
            execution_time = time.time() - start_time
            logger.info(f"Transaction execution time: {execution_time:.4f} seconds")

            if tx_status.success:
                logger.info(f"Front-run successful: {tx_status.signature}")
                
                # Create position tracking data
                self.current_position = {
                    "entry_sig": tx_status.signature,
                    "victim_sig": victim_signature,
                    "token_mint": token_mint,
                    "entry_time": time.time(),
                    "entry_sol": FRONT_RUN_AMOUNT
                }
                
                # Schedule exit after a short delay to capture profit
                asyncio.create_task(self.schedule_exit())
            else:
                logger.error(f"Front-run failed: {tx_status.error}")
                self.stats['failure_count'] += 1

        except Exception as e:
            logger.error(f"Execution error: {e}")
            logger.debug(traceback.format_exc())
            self.stats['failure_count'] += 1
        finally:
            # Remove from active trades after completion
            self.active_trades.discard(victim_signature)
            
    async def schedule_exit(self):
        """Schedule position exit after delay with optimized timing"""
        try:
            # Calculate optimal exit delay based on victim tx processing time
            delay = EXIT_DELAY
            logger.info(f"Scheduling exit in {delay:.2f} seconds")
            
            # Wait for victim transaction to complete
            await asyncio.sleep(delay)
            
            # Attempt exit with retries
            max_attempts = 3
            for attempt in range(1, max_attempts + 1):
                exit_success = await self.exit_position()
                if exit_success:
                    break
                
                if attempt < max_attempts:
                    retry_delay = attempt * 2  # Exponential backoff
                    logger.warning(f"Exit attempt {attempt} failed, retrying in {retry_delay} seconds...")
                    await asyncio.sleep(retry_delay)
            
            # If still have position after max attempts, log the issue
            if attempt == max_attempts and self.current_position:
                logger.error("Failed to exit position after maximum attempts")
                
        except Exception as e:
            logger.error(f"Error in schedule_exit: {e}")
            logger.debug(traceback.format_exc())
            
    async def exit_position(self):
        """Exit current market position with optimized transaction handling"""
        if not self.current_position:
            return True  # No position to exit

        try:
            token_mint = self.current_position["token_mint"]
            logger.info(f"Exiting position for {token_mint}")
            
            # Get token balance
            balances = await self.trading_core.get_balances()
            _, token_balances = balances
            token_balance = token_balances.get(token_mint, 0)
            
            # Skip if balance is too small
            min_exit_balance = 0.00001
            if token_balance <= min_exit_balance:
                logger.warning(f"Token balance too small to exit: {token_balance}")
                self.current_position = None
                return True
                
            logger.info(f"Selling token balance: {token_balance}")
            
            # Calculate position duration
            entry_time = self.current_position.get("entry_time", 0)
            position_duration = time.time() - entry_time
            logger.info(f"Position held for {position_duration:.2f} seconds")
            
            # Create swap parameters for the exit (using 98% of balance to ensure success)
            sell_params = SwapParams(
                token_mint=token_mint,
                amount_in=token_balance * 0.98,  # Small buffer to ensure success
                min_amount_out=None,  # No minimum protection, prioritize exit
                is_buy=False,  # This is a sell
                slippage_bps=SLIPPAGE_BPS
            )
            
            # Execute the exit swap
            tx_status = await self.trading_core.execute_swap(sell_params)

            if tx_status.success:
                logger.info(f"Exit successful: {tx_status.signature}")
                
                # Calculate profit
                entry_sig = self.current_position.get("entry_sig", "")
                
                # Get updated balances
                if self.balance_monitor:
                    # Use WebSocket-based balance monitoring (more accurate)
                    sol_balance = self.balance_monitor.sol_balance
                    profit = sol_balance - (self.stats.get('last_sol_balance', sol_balance) - FRONT_RUN_AMOUNT)
                    self.stats['last_sol_balance'] = sol_balance
                else:
                    # Fallback to RPC balance check
                    updated_balances = await self.trading_core.get_balances()
                    sol_balance, _ = updated_balances
                    profit = sol_balance - (self.stats.get('last_sol_balance', sol_balance) - FRONT_RUN_AMOUNT)
                    self.stats['last_sol_balance'] = sol_balance
                
                logger.info(f"Trade complete: {entry_sig} → {tx_status.signature}")
                logger.info(f"Estimated profit: {profit:.6f} SOL")
                
                # Update stats
                self.stats['last_profit'] = profit
                self.stats['profit_total'] += profit
                self.stats['success_count'] += 1
                
                # Clear current position
                self.current_position = None
                return True
            else:
                logger.error(f"Exit failed: {tx_status.error}")
                return False

        except Exception as e:
            logger.error(f"Exit error: {e}")
            logger.debug(traceback.format_exc())
            return False
    
    async def run(self):
        """Main bot execution loop with proper lifecycle management"""
        logger.info("Starting Ultimate Raydium MEV Bot with GPU acceleration")
        logger.info(f"Trading parameters: FRONT_RUN_AMOUNT={FRONT_RUN_AMOUNT}, MIN_PROFIT={MIN_PROFIT}")
        
        try:
            # Initialize the bot
            success = await self.initialize()
            if not success:
                logger.error("Bot initialization failed")
                return
            
            # Start monitoring the mempool
            await self.monitor_mempool()
            
        except KeyboardInterrupt:
            logger.info("Bot stopped by user")
        except Exception as e:
            logger.error(f"Fatal error: {e}")
            logger.debug(traceback.format_exc())
        finally:
            # Ensure proper cleanup
            await self.shutdown()
            
    async def shutdown(self):
        """Handle shutdown with proper resource cleanup"""
        if not self.is_running:
            return  # Already shutting down
            
        logger.info("Shutting down bot...")
        self.is_running = False
        
        # Exit any current position before shutting down
        if self.current_position:
            logger.info("Exiting current position before shutdown")
            await self.exit_position()
            
        # Stop balance monitor
        if self.balance_monitor:
            await self.balance_monitor.stop()
            
        # Release GPU resources
        if hasattr(self, 'gpu'):
            self.gpu.cleanup()
            
        # Clean up thread pool
        self.thread_pool.shutdown(wait=True)
            
        logger.info("Shutdown complete")
        
        # Set shutdown event
        self.shutdown_event.set()

#==============================================================================
# Main Entrypoint
#==============================================================================

async def main():
    """Main entry point with proper startup and shutdown handling"""
    # Display banner
    print("""
    ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
    ┃                Ultimate Raydium MEV Bot v1.0 - NVIDIA Edition         ┃
    ┃               Optimized for 4070 Ti with CUDA 11.8 & PyTorch          ┃
    ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
    """)
    
    # Initialize the bot
    bot = RaydiumMEVBot()
    
    try:
        # Run the bot
        await bot.run()
        
        # Wait for shutdown signal
        await bot.shutdown_event.wait()
        
    except KeyboardInterrupt:
        logger.info("Bot stopped by user")
        await bot.shutdown()
    except Exception as e:
        logger.error(f"Fatal error: {e}")
        logger.error(traceback.format_exc())
        await bot.shutdown()
        sys.exit(1)

if __name__ == "__main__":
    try:
        # Use uvloop for better performance if available
        try:
            import uvloop
            uvloop.install()
            logger.info("uvloop installed for improved performance")
        except ImportError:
            .sol_balance - self.balance_monitor.last_sol_balance
                    
                    if sol_change > 0:
                        # Calculate profit based on SOL balance increase
                        profit = sol_change
                        logger.info(f"Estimated profit: {profit:.6f} SOL")
                        
                        # Update stats
                        self.stats['last_profit'] = profit
                        self.stats['profit_total'] += profit
                        self.stats['success_count'] += 1
                        
                        # Clear current position
                        self.current_position = None
                        
            # General SOL balance increase might indicate profit from other trades
            if balance_type == 'SOL' and new_value > old_value:
                # Only record significant balance changes
                change = new_value - old_value
                if change > 0.001 and not self.current_position:
                    logger.info(f"Detected SOL balance increase: +{change:.6f} SOL")
                    
                    # Record this as potential profit if not associated with a specific trade
                    if time.time() - self.last_trade_time > 10:
                        self.stats['profit_total'] += change
                        logger.info(f"Added {change:.6f} SOL to total profit")
                        
        except Exception as e:
            logger.error(f"Error handling balance change: {e}")
    
    async def healthcheck(self):
        """Periodic healthcheck to monitor system status"""
        try:
            while self.is_running:
                # Calculate uptime
                uptime = time.time() - self.stats['start_time']
                hours = uptime / 3600
                
                # Get current balances
                sol_balance = self.balance_monitor.sol_balance if self.balance_monitor else 0
                wsol_balance = self.balance_monitor.wsol_balance if self.balance_monitor else 0
                
                logger.info(
                    f"Bot status: Running for {hours:.2f}h | "
                    f"SOL: {sol_balance:.4f} | WSOL: {wsol_balance:.4f} | "
                    f"Active trades: {len(self.active_trades)} | "
                    f"Position: {'Open' if self.current_position else 'None'} | "
                    f"Success/Fail: {self.stats['success_count']}/{self.stats['failure_count']} | "
                    f"Profit: {self.stats['profit_total']:.6f} SOL"
                )
                
                # Check GPU health if using CUDA
                if torch.cuda.is_available():
                    mem_used = torch.cuda.memory_allocated() / 1024 / 1024 / 1024
                    mem_reserved = torch.cuda.memory_reserved() / 1024 / 1024 / 1024
                    logger.debug(f"GPU memory: {mem_used:.2f} GB used, {mem_reserved:.2f} GB reserved")
                
                # Reset rate limiting if needed
                now = time.time()
                if now - self.minute_start > 60:
                    self.trades_this_minute = 0
                    self.minute_start = now
                
                await asyncio.sleep(30)
                
        except asyncio.CancelledError:
            logger.info("Healthcheck task cancelled")
        except Exception as e:
            logger.error(f"Healthcheck error: {e}")
            logger.error(traceback.format_exc())

    async def monitor_mempool(self):
        """Monitor Solana mempool for Raydium swap transactions"""
        logger.info(f"Starting mempool monitoring with connection to {HELIUS_WS_URL}")
        
        # Use Helius WebSocket for optimized performance
        ws_url = HELIUS_WS_URL
        
        # Configure headers with API key
        headers = {"api-key": API_KEY} if API_KEY else {}
        headers["User-Agent"] = "solana-client/1.18.23 (client:SolanaLabs)"
        
        reconnect_attempts = 0
        max_reconnect_delay = 60  # Maximum backoff in seconds
        
        while self.is_running:
            try:
                # Connect to WebSocket
                async with websockets.connect(ws_url, extra_headers=headers) as ws:
                    logger.info("Connected to mempool WebSocket")
                    reconnect_attempts = 0
                    
                    # Subscribe to Raydium program logs
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "logsSubscribe",
                        "params": [
                            {"mentions": [RAYDIUM_PROGRAM_ID]},
                            {"commitment": "processed"}  # Fastest commitment level
                        ]
                    }))
                    
                    # Wait for subscription confirmation
                    response = await asyncio.wait_for(ws.recv(), timeout=10)
                    data = json.loads(response)
                    
                    if "result" in data:
                        sub_id = data["result"]
                        logger.info(f"Mempool subscription established: ID {sub_id}")
                    else:
                        logger.error(f"Failed to establish mempool subscription: {data}")
                        if "error" in data:
                            logger.error(f"Error: {data['error']}")
                        continue
                    
                    # Process incoming messages
                    while self.is_running:
                        try:
                            # Wait for message with timeout
                            msg = await asyncio.wait_for(ws.recv(), timeout=30)
                            await self.process_transaction(msg)
                        except asyncio.TimeoutError:
                            # No message received, send ping to keep connection alive
                            await ws.ping()
                            continue
            
            except (websockets.exceptions.ConnectionClosed, 
                    websockets.exceptions.WebSocketException) as e:
                if not self.is_running:
                    break
                    
                # Implement backoff for reconnection attempts
                reconnect_attempts += 1
                backoff = min(5 * reconnect_attempts, max_reconnect_delay)
                logger.error(f"Mempool WebSocket error: {e}. "
                           f"Reconnecting in {backoff} seconds...")
                await asyncio.sleep(backoff)
                
            except Exception as e:
                if not self.is_running:
                    break
                    
                logger.error(f"Unexpected error in mempool monitor: {e}")
                logger.error(traceback.format_exc())
                await asyncio.sleep(5)

    async def process_transaction(self, msg: str):
        """Process incoming transaction data with GPU-accelerated analysis"""
        try:
            data = json.loads(msg)
            
            # Check if this is a proper logs notification
            if "method" not in data or data["method"] != "logsNotification":
                return
                
            # Extract logs and signature
            logs = data.get('params', {}).get('result', {}).get('value', {}).get('logs', [])
            signature = data.get('params', {}).get('result', {}).get('value', {}).get('signature', '')
            
            # Skip if no logs or already processed this transaction
            if not logs or signature in self.seen_txs:
                return
            
            # Add to seen transactions
            self.seen_txs.add(signature)
            
            # Check if this is our own transaction
            if signature in self.active_trades:
                logger.info(f"Received confirmation for our transaction: {signature}")
                return
                
            # Skip if we already have a position
            if self.current_position:
                return
                
            # Enforce rate limiting
            if self.trades_this_minute >= MAX_TRADES_PER_MINUTE:
                return
                
            # Quickly check for Raydium program and swap instruction
            has_raydium_program = False
            has_swap_instruction = False
            
            for log in logs:
                if f"Program {RAYDIUM_PROGRAM_ID} invoke" in log:
                    has_raydium_program = True
                if "Program log: Instruction: Swap" in log:
                    has_swap_instruction = True
                    break
            
            # If potential swap found, extract details and check profitability
            if has_raydium_program and has_swap_instruction:
                # Extract swap details with optimized parsing
                swap_data = await self.parse_swap_instruction(logs, signature)
                
                if swap_data:
                    # Check profitability asynchronously with a short timeout
                    is_profitable_task = asyncio.create_task(self.is_profitable(swap_data))
                    try:
                        profitable = await asyncio.wait_for(is_profitable_task, timeout=0.05)
                        if profitable:
                            # Execute front-run if profitable
                            await self.execute_front_run(swap_data, signature)
                    except asyncio.TimeoutError:
                        # If profitability check takes too long, move on
                        is_profitable_task.cancel()

        except json.JSONDecodeError:
            # Not valid JSON, ignore silently
            pass
        except Exception as e:
            logger.error(f"Transaction processing error: {e}")
            logger.debug(traceback.format_exc())

    async def parse_swap_instruction(self, logs: List[str], signature: str) -> Optional[Dict]:
        """Extract swap details from transaction logs with optimized parsing"""
        try:
            # Fast path: extract key details with minimal processing
            token_mint = None
            amount_in_sol = None
            pool_id = None
            
            # Track instruction details
            is_raydium_v4 = False
            has_swap_instruction = False
            has_instruction_9 = False
            
            # Extract instruction data (loop through logs once)
            for log in logs:
                # Check for Raydium program invoke (V4)
                if f"Program {RAYDIUM_PROGRAM_ID} invoke" in log:
                    is_raydium_v4 = True
                    continue
                    
                # Check for explicit swap instruction
                if "Program log: Instruction: Swap" in log and is_raydium_v4:
                    has_swap_instruction = True
                    continue
                    
                # Check for AMM ID (pool identifier)
                if is_raydium_v4 and "Pool" in log and len(log.strip()) >= 32:
                    potential_pool = log.strip()
                    if all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in potential_pool):
                        pool_id = potential_pool
                    continue
                    
                # Extract amount from instruction data
                if "Program data: " in log and is_raydium_v4:
                    try:
                        data = base64.b64decode(log.split("Program data: ")[1])
                        if len(data) > 0 and data[0] == RAYDIUM_SWAP_INSTRUCTION:  # Check for 0x09
                            has_instruction_9 = True
                            
                            # Extract amount from first 8 bytes after instruction type
                            if len(data) >= 9:
                                amount_lamports = int.from_bytes(data[1:9], byteorder='little')
                                amount_in_sol = amount_lamports / LAMPORTS_PER_SOL
                                logger.debug(f"Extracted amount: {amount_in_sol:.4f} SOL")
                    except Exception as e:
                        logger.debug(f"Error extracting instruction data: {e}")
                    continue
                
                # Look for token mint in token program logs
                if token_mint is None and TOKEN_PROGRAM_ID in log:
                    parts = log.split()
                    for i, part in enumerate(parts):
                        # Valid token mint addresses are base58 encoded and 32+ chars
                        if len(part) >= 32 and all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in part):
                            # Skip if it's our wallet or WSOL (might be false positive)
                            if part != self.wallet_pubkey and part != WSOL_ADDRESS:
                                token_mint = part
                                break
            
            # Skip if not a valid V4 AMM transaction with instruction 9
            if not is_raydium_v4 or not has_instruction_9:
                return None
                
            # If we couldn't find a token, try harder (sometimes embedded in other logs)
            if token_mint is None:
                for log in logs:
                    if "token:" in log.lower() or "Token:" in log:
                        parts = log.split()
                        for part in parts:
                            if len(part) >= 32 and all(c in '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz' for c in part):
                                if part != self.wallet_pubkey and part != WSOL_ADDRESS:
                                    token_mint = part
                                    break
                    if token_mint:
                        break
            
            # Skip if we couldn't identify a token
            if not token_mint:
                return None
                
            # Normalize SOL address to WSOL if needed (handle system program address)
            if token_mint == SYSTEM_PROGRAM_ID:
                token_mint = WSOL_ADDRESS
                
            # Use default amount if we couldn't extract it
            if not amount_in_sol:
                amount_in_sol = 1.0  # Conservative default
                
            # Create swap data structure with all extracted information
            swap_data = {
                "token_mint": token_mint,
                "amount_in_sol": amount_in_sol,
                "signature": signature,
                "direction": "buy",  # Assume buying tokens with SOL
                "timestamp": time.time()
            }
            
            # Add pool ID if found (useful for multi-pool tokens)
            if pool_id:
                swap_data["pool_id"] = pool_id
                
            logger.debug(f"Found potential swap: {token_mint}, {amount_in_sol:.4f} SOL")
            return swap_data
            
        except Exception as e:
            logger.error(f"Error parsing swap instruction: {e}")
            logger.debug(traceback.format_exc())
            return None
            
    async def is_profitable(self, swap_data: Dict) -> bool:
        """Determine if a front-running opportunity is profitable"""
        try:
            # Extract key data
            token_mint = swap_data.get("token_mint")
            victim_amount = swap_data.get("amount_in_sol", 1.0)
            
            # Skip transactions with missing data
            if not token_mint:
                return False
                
            # Check if transaction is too recent (prevent double-front-running)
            now = time.time()
            if now - self.last_trade_time < 0.5:  # 500ms cooldown
                logger.debug("Transaction too soon after last trade")
                return False
                
            # Check token safety
            # Use cached result if available
            current_time = time.time()
            if token_mint in self.token_safety_cache:
                cache_time, safety = self.token_safety_cache[token_mint]
                if current_time - cache_time < SECURITY_CACHE_SECONDS:
                    if not safety.is_safe:
                        logger.debug(f"Token {token_mint} failed cached safety check")
                        return False
                else: DONT TAKE AWAY MY IMPORTS I WANT TO SEE ALL MY IMPORTS ALL MY CODE JUST REMOVE ANY DUPLICATE CODE AND REORGANIZE THE SPACING PROPERLY PLEASE