I need the exact word for word script back with nothing changed or removed other then proper spacing got it #!/usr/bin/env python3
"""
Quantum Raydium MEV Bot - NO VALIDATION VERSION
----------------------------------------------------
High-performance MEV trader with GPU acceleration and
immediate trade execution. NO TOKEN VALIDATION - will trade any token.
"""

import asyncio
import base58
import base64
import json
import logging
import os
import signal
import sys
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Any
import re
import ssl
import aiohttp
import websockets
from dotenv import load_dotenv

# Import GPU acceleration
try:
    import torch
    import numpy as np
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    logging.warning("PyTorch not installed. Running without GPU acceleration.")

# Configure logging
os.makedirs("logs", exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    handlers=[
        logging.FileHandler(f"logs/quantum_mev_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Fix for Windows event loop
if sys.platform == 'win32':
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

# Import Rust trading core
from solana_rust_bot import SolanaTrader

# Constants
LAMPORTS_PER_SOL = 1_000_000_000

@dataclass
class TokenSecurity:
    mint_authority_revoked: bool = True  # DEFAULT TO TRUE - BYPASS VALIDATION
    lp_locked: bool = True               # DEFAULT TO TRUE - BYPASS VALIDATION
    has_exit_liquidity: bool = True      # DEFAULT TO TRUE - BYPASS VALIDATION
    exit_liquidity_sol: float = 50.0     # DEFAULT TO HIGH - BYPASS VALIDATION
    pool_address: str = "BYPASS_VALIDATION"
    token_mint: str = ""
    valid: bool = True                   # DEFAULT TO TRUE - BYPASS VALIDATION
    exists: bool = True                  # DEFAULT TO TRUE - BYPASS VALIDATION
    
@dataclass
class TradingStats:
    transactions_seen: int = 0
    raydium_swaps_found: int = 0
    viable_opportunities: int = 0
    trades_attempted: int = 0
    trades_completed: int = 0
    failed_front_runs: int = 0
    failed_back_runs: int = 0
    total_profit: float = 0.0
    total_loss: float = 0.0
    start_time: float = field(default_factory=time.time)
    token_profits: Dict[str, float] = field(default_factory=dict)
    victim_confirmed_count: int = 0
    victim_missed_count: int = 0
    initial_sol_balance: float = 0.0
    rpc_calls: int = 0
    cached_tokens_used: int = 0
    ws_connections: int = 0
    ws_messages_received: int = 0
    ws_errors: int = 0
    ws_reconnects: int = 0
    invalid_tokens_count: int = 0  
    account_not_found_count: int = 0
    whitelist_hits: int = 0
    gpu_processing_time: float = 0.0
    
@dataclass
class TradeOpportunity:
    token_mint: str
    amount_sol: float
    tx_signature: str
    accounts: Dict[str, str]
    front_amount: float
    start_time: float = field(default_factory=time.time)
    front_tx: Optional[str] = None
    back_tx: Optional[str] = None
    status: str = "pending"
    profit: float = 0.0
    victim_confirmed: bool = False
    token_balance: float = 0.0
    token_decimals: int = 9
    start_sol_balance: float = 0.0
    start_wsol_balance: float = 0.0
    end_sol_balance: float = 0.0
    end_wsol_balance: float = 0.0

class RaydiumDataExtractor:
    """Extract and validate token data from Raydium pools and transaction logs"""

    def __init__(self, tensor_size=176):
        self.tensor_size = tensor_size
        
        # Initialize GPU acceleration if available
        if HAS_TORCH and torch.cuda.is_available():
            self.device = torch.device('cuda')
            # Setup CUDA for maximum performance
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False
            # Create stream for parallel processing
            self.stream = torch.cuda.Stream(priority=-3)
            self.warmup_gpu()
            logger.info(f"GPU acceleration enabled: {torch.cuda.get_device_name()}")
        else:
            self.device = None
            logger.info("Running without GPU acceleration")
        
        # Program IDs to exclude when identifying token mints
        self.program_ids = {
            os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8"),
            os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
            os.getenv("SERUM_PROGRAM_ID", "9xQeWvG816bUx9EPjHmaT23yvVM2ZWbrrpZb9PusVFin"),
            os.getenv("SYSTEM_PROGRAM_ID", "11111111111111111111111111111111"),
            os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112"),
            os.getenv("ASSOCIATED_TOKEN_PROGRAM_ID", "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"),
        }
        
    def warmup_gpu(self):
        """Warm up GPU to ensure consistent performance"""
        if not HAS_TORCH or not torch.cuda.is_available():
            return
            
        # Create tensors for warmup
        with torch.cuda.stream(self.stream):
            warmup_tensor = torch.randn(self.tensor_size, device=self.device)
            for _ in range(int(os.getenv("WARMUP_CYCLES", "2500"))):  # Warm up cycles from benchmarks
                result = torch.matmul(warmup_tensor, warmup_tensor)
                
        # Synchronize to ensure warmup is complete
        torch.cuda.synchronize()
        logger.info("GPU warmup completed")
    
    def accelerate_log_parsing(self, logs):
        """Use GPU acceleration for log parsing"""
        if not HAS_TORCH or not torch.cuda.is_available() or not self.device:
            # Fallback to CPU
            return self.extract_token_mint(logs)
            
        start_time = time.time()
        try:
            # If using GPU, we can accelerate pattern matching with tensor operations
            log_combined = " ".join(logs)
            
            # Extract all potential Solana addresses
            address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
            matches = re.findall(address_pattern, log_combined)
            
            if not matches:
                return None
                
            # Filter known program IDs
            filtered_matches = [m for m in matches if m not in self.program_ids]
            
            if not filtered_matches:
                return None
                
            # Convert addresses to tensors for faster processing
            with torch.cuda.stream(self.stream):
                # Create encoding for each address (use first 8 chars as feature)
                addr_features = []
                for addr in filtered_matches:
                    # Simple feature: sum of ASCII values of first 8 chars
                    feature = sum(ord(c) for c in addr[:8])
                    addr_features.append(feature)
                
                # Convert to tensor 
                if addr_features:
                    features_tensor = torch.tensor(addr_features, device=self.device, dtype=torch.float32)
                    
                    # Calculate frequency score based on occurrence pattern
                    scores = []
                    for addr in filtered_matches:
                        # Count occurrences
                        score = log_combined.count(addr) * 10
                        
                        # Bonus for appearing in ray_log lines
                        for log in logs:
                            if "ray_log:" in log and addr in log:
                                score += 50
                        
                        # Bonus for appearing with "token" or "mint"
                        for log in logs:
                            if ("token" in log.lower() or "mint" in log.lower()) and addr in log:
                                score += 30
                                
                        scores.append(score)
                    
                    scores_tensor = torch.tensor(scores, device=self.device, dtype=torch.float32)
                    
                    # Get highest score
                    max_score, max_idx = torch.max(scores_tensor, dim=0)
                    
                    # Get the address with highest score
                    token_mint = filtered_matches[max_idx.item()]
                    
                    # Verify it's a valid pubkey
                    if self.is_valid_pubkey(token_mint):
                        torch.cuda.synchronize()
                        processing_time = time.time() - start_time
                        logger.debug(f"GPU token extraction completed in {processing_time*1000:.2f}ms")
                        return token_mint
            
            # Fallback to CPU method if GPU processing fails
            return self.extract_token_mint(logs)
                
        except Exception as e:
            logger.error(f"Error in GPU acceleration: {e}")
            # Fallback to CPU method
            return self.extract_token_mint(logs)
        
    def decode_ray_log(self, log_str):
        """Decode Raydium ray_log with optimized parsing"""
        if "ray_log:" not in log_str:
            return None
            
        try:
            ray_log_part = log_str.split("ray_log:")[1].strip()
            
            # Multi-strategy decoding with validation
            raw_data = None
            decode_attempts = [
                lambda x: base64.b64decode(x),  # Standard
                lambda x: base64.b64decode(x + "=" * (4 - len(x) % 4) % 4),  # Add padding
                lambda x: base64.b64decode(x.replace("-", "+").replace("_", "/")),  # URL-safe
                lambda x: base64.b64decode(x.replace("-", "+").replace("_", "/") + "=" * (4 - len(x.replace("-", "+").replace("_", "/")) % 4) % 4)  # URL-safe with padding
            ]
            
            # Try all decoding strategies
            for decode_func in decode_attempts:
                try:
                    raw_data = decode_func(ray_log_part)
                    if raw_data and len(raw_data) >= 9:
                        break
                except Exception:
                    continue
                    
            if not raw_data or len(raw_data) < 9:
                return None
            
            # First byte is instruction type
            instruction_type = raw_data[0]
            
            # Next 8 bytes are the amount in little-endian
            amount_bytes = raw_data[1:9]
            amount = int.from_bytes(amount_bytes, byteorder="little") / LAMPORTS_PER_SOL
            
            # Consider both 0x03 and 0x09 as valid buy indicators
            is_buy = instruction_type in [3, 9]  # 3 & 9 are buy instructions
            is_sell = instruction_type == 11  # 11 is sell instruction
            
            # Extract pubkeys if available
            pubkeys = []
            if len(raw_data) > 9:
                account_data = raw_data[9:]
                
                for i in range(0, len(account_data) - 31, 32):
                    try:
                        pubkey_bytes = account_data[i:i+32]
                        # Skip zero or all-one pubkeys
                        if all(b == 0 for b in pubkey_bytes) or all(b == 255 for b in pubkey_bytes):
                            continue
                            
                        pubkey = base58.b58encode(pubkey_bytes).decode('utf-8')
                        if self.is_valid_pubkey(pubkey):
                            pubkeys.append(pubkey)
                    except Exception:
                        continue
            
            return {
                "instruction_type": instruction_type,
                "amount": amount,
                "is_buy": is_buy,
                "is_sell": is_sell,
                "pubkeys": pubkeys,
                "raw_data": raw_data
            }
            
        except Exception as e:
            logger.error(f"Error decoding ray_log: {e}")
            return None
            
    def is_valid_pubkey(self, address):
        """Validate a Solana address format"""
        if not address or len(address) < 32 or len(address) > 44:
            return False
            
        try:
            # Check for concatenated addresses (e.g. 11111111... + another address)
            if address.startswith("11111111") and len(address) > 8:
                return False
                
            # Make sure we decode to 32 bytes exactly
            decoded = base58.b58decode(address)
            return len(decoded) == 32
        except Exception:
            return False
    
    def extract_token_mint(self, logs):
        """Extract token mint from logs with multi-pass approach and better validation"""
        # Method 1: Extract from Transfer logs targeting Mint Authority
        token_program_id = os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
        token_candidates = []
        
        for log in logs:
            if token_program_id in log and "Instruction: Transfer" in log:
                # Look for the next few logs for pubkeys
                token_index = logs.index(log)
                for i in range(token_index, min(token_index + 5, len(logs))):
                    address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                    matches = re.findall(address_pattern, logs[i])
                    for match in matches:
                        if match not in self.program_ids and self.is_valid_pubkey(match):
                            token_candidates.append(match)
        
        # Method 2: Extract from ray_log pubkeys
        for log in logs:
            if "ray_log:" in log:
                ray_data = self.decode_ray_log(log)
                if ray_data and ray_data["pubkeys"]:
                    for pubkey in ray_data["pubkeys"]:
                        # Skip known program IDs
                        if pubkey not in self.program_ids and self.is_valid_pubkey(pubkey):
                            token_candidates.append(pubkey)
        
        # Method 3: Extract from MintTo or related token instructions
        for log in logs:
            if token_program_id in log and "Program data:" in log and "MintTo" in log:
                address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                matches = re.findall(address_pattern, log)
                for match in matches:
                    if match not in self.program_ids and self.is_valid_pubkey(match):
                        token_candidates.append(match)
        
        # Count occurrences and pick the most likely token mint
        if token_candidates:
            # Filter out common program IDs and system accounts
            filtered_candidates = [c for c in token_candidates if not c.startswith("1111")]
            
            if filtered_candidates:
                # Count occurrences and pick the most frequent
                token_counts = {}
                for token in filtered_candidates:
                    token_counts[token] = token_counts.get(token, 0) + 1
                
                # Sort by frequency and return the most common
                sorted_tokens = sorted(token_counts.items(), key=lambda x: x[1], reverse=True)
                return sorted_tokens[0][0]
        
        return None
        
    def extract_accounts(self, logs):
        """Extract account information from logs"""
        accounts = {
            "token_program_id": os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
            "serum_program_id": os.getenv("SERUM_PROGRAM_ID", "9xQeWvG816bUx9EPjHmaT23yvVM2ZWbrrpZb9PusVFin"),
            "raydium_program_id": os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8"),
            "system_program_id": os.getenv("SYSTEM_PROGRAM_ID", "11111111111111111111111111111111")
        }
        
        # Account field mapping from logs
        account_fields = {
            "amm id": "amm_id",
            "authority": "amm_authority",
            "open orders": "amm_open_orders",
            "target orders": "amm_target_orders",
            "pool coin": "pool_coin_token_account",
            "pool pc": "pool_pc_token_account",
            "serum market": "serum_market",
            "serum bids": "serum_bids",
            "serum asks": "serum_asks",
            "event queue": "serum_event_queue",
            "coin vault": "serum_coin_vault",
            "pc vault": "serum_pc_vault",
            "vault signer": "serum_vault_signer"
        }
        
        # Use GPU acceleration for account extraction if available
        if HAS_TORCH and torch.cuda.is_available() and hasattr(self, 'device') and self.device:
            try:
                with torch.cuda.stream(self.stream):
                    # Prepare tensor of all logs
                    log_data = " ".join(logs)
                    
                    # Process multiple matches in parallel
                    for pattern, field_name in account_fields.items():
                        if pattern in log_data.lower():
                            # Find the log line containing this pattern
                            for log in logs:
                                if pattern in log.lower() and ":" in log:
                                    parts = log.split(":")
                                    if len(parts) < 2:
                                        continue
                                        
                                    # Extract pubkey after pattern
                                    address_part = parts[1].strip()
                                    address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                                    matches = re.findall(address_pattern, address_part)
                                    
                                    for match in matches:
                                        if self.is_valid_pubkey(match):
                                            accounts[field_name] = match
                                            break
                    
                # Synchronize GPU operations
                torch.cuda.synchronize()
                
            except Exception as e:
                logger.error(f"GPU acceleration failed for account extraction: {e}")
                # Continue with CPU-based approach below
        
        # CPU-based approach (standard)
        if len(accounts) <= 4:  # If GPU didn't find accounts, use CPU approach
            for log in logs:
                if ":" not in log:
                    continue
                    
                # Check for account patterns
                for pattern, field_name in account_fields.items():
                    if pattern in log.lower():
                        parts = log.split(":")
                        if len(parts) < 2:
                            continue
                            
                        # Extract pubkey after pattern
                        address_part = parts[1].strip()
                        address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                        matches = re.findall(address_pattern, address_part)
                        
                        for match in matches:
                            if self.is_valid_pubkey(match):
                                accounts[field_name] = match
                                break
        
        # Extract AMM ID from ray_log as fallback
        if "amm_id" not in accounts:
            for log in logs:
                if "ray_log:" in log:
                    ray_data = self.decode_ray_log(log)
                    if ray_data and ray_data["pubkeys"] and len(ray_data["pubkeys"]) >= 5:
                        accounts["amm_id"] = ray_data["pubkeys"][0]
                        break
        
        return accounts

class GPUAcceleratedTokenValidator:
    """BYPASSED Token validator - accepts all tokens"""
    
    def __init__(self, trader, rpc_url):
        self.trader = trader
        self.rpc_url = rpc_url
        self.security_cache = {}
        
        # Initialize GPU acceleration if available
        if HAS_TORCH and torch.cuda.is_available():
            self.device = torch.device('cuda')
            # Setup CUDA for maximum performance
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False
            # Allocate memory fraction for validation
            gpu_memory_fraction = float(os.getenv("GPU_MEMORY_FRACTION", "0.95"))
            if hasattr(torch.cuda, 'set_per_process_memory_fraction'):
                torch.cuda.set_per_process_memory_fraction(gpu_memory_fraction)
            logger.info(f"GPU acceleration enabled for validation: {torch.cuda.get_device_name()}")
        else:
            self.device = None
            logger.info("Running token validation without GPU acceleration")
        
    async def validate_token_security(self, token_mint):
        """BYPASSED: Accept any token without validation"""
        # Always return a valid token security object
        security = TokenSecurity(
            mint_authority_revoked=True,  # Force to True
            lp_locked=True,               # Force to True
            has_exit_liquidity=True,      # Force to True
            exit_liquidity_sol=50.0,      # Force high liquidity
            pool_address="BYPASS_VALIDATION",
            token_mint=token_mint,
            valid=True,                   # Force to True
            exists=True                   # Force to True
        )
        
        # Cache result to avoid RPC calls
        self.security_cache[token_mint] = (time.time(), security)
        
        logger.info(f"BYPASS: Accepting token {token_mint[:8]} without validation")
        return security
    
    # Keeping these methods as stubs to avoid breaking code elsewhere
    async def verify_token_exists(self, token_mint):
        return True
        
    async def get_token_pool(self, token_mint):
        return "BYPASS_VALIDATION"
    
    async def get_pool_data(self, pool_address):
        return bytes([0] * 680)  # Return dummy data
    
    def analyze_pool_security(self, data):
        return {
            "exit_liquidity": 50.0,
            "lp_locked": True
        }
    
    def verify_security_features(self, data):
        return (True, True)  # Always return (mint_revoked, lp_locked) as both True

class QuantumMEVTrader:
    """High-performance MEV trader for Raydium with GPU acceleration and token security validation"""
    
    def __init__(self):
        # Load environment variables
        load_dotenv()
        
        # Create thread pool for parallel processing
        self.thread_pool = ThreadPoolExecutor(
            max_workers=int(os.getenv("THREADS", "48"))
        )
        
        # Initialize trader with keypair
        keypair_path = os.getenv("SOLANA_KEYPAIR", os.getenv("KEYPAIR_PATH", "C:\\solana_rust_bot\\keypair.bin"))
        
        self.trader = SolanaTrader(
            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"),
            keypair_path=keypair_path,
            compute_limit=int(os.getenv("COMPUTE_LIMIT", "50000"))
        )
        
        # Initialize wallet info
        self.wallet_address = self.trader.get_address()
        
        # Initialize components with GPU acceleration
        self.token_extractor = RaydiumDataExtractor(
            tensor_size=int(os.getenv("TENSOR_SIZE", "176"))
        )
        self.token_validator = GPUAcceleratedTokenValidator(
            self.trader, 
            os.getenv("RPC_URL", "https://winny-rychu7-fast-mainnet.helius-rpc.com")
        )
        
        # State tracking
        self.running = False
        self.active_opportunity = None
        self.processed_txs = set()
        self.valid_tokens = set()
        self.invalid_tokens = set()
        self.last_balance_update = 0
        self.sol_balance = 0
        self.wsol_balance = 0
        self.wsol_account = os.getenv("WSOL_TOKEN_ACCOUNT", "")
        self.stats = TradingStats()
        self.balance_lock = asyncio.Lock()
        
        # Config
        self.front_run_amount = float(os.getenv("FRONT_RUN_AMOUNT", "0.05"))
        self.min_victim_size = float(os.getenv("MIN_VICTIM_SIZE", "5.0"))
        self.max_victim_size = float(os.getenv("MAX_VICTIM_SIZE", "320.0"))
        self.min_sol_balance = float(os.getenv("MIN_SOL_BALANCE", "0.03"))
        self.min_wsol_balance = float(os.getenv("MIN_WSOL_BALANCE", "0.02"))
        self.slippage_bps = int(os.getenv("SLIPPAGE_BPS", 1500))
        self.min_profit = float(os.getenv("MIN_PROFIT", "0.001"))
        self.wsol_wrap_interval = int(os.getenv("WSOL_WRAP_INTERVAL", "30"))
        self.last_wrap_time = 0
        
        # Load whitelist if available
        self.whitelist_path = "data/valid_tokens.json"
        self.load_whitelist()
        
        logger.info(f"Initialized Quantum MEV trader with wallet: {self.wallet_address}")
    
    def load_whitelist(self):
        """Load whitelist from file if it exists"""
        try:
            os.makedirs("data", exist_ok=True)
            if os.path.exists(self.whitelist_path):
                with open(self.whitelist_path, 'r') as f:
                    data = json.load(f)
                    self.valid_tokens = set(data.get("valid_tokens", []))
                    self.invalid_tokens = set(data.get("invalid_tokens", []))
                logger.info(f"Loaded whitelist with {len(self.valid_tokens)} valid tokens and {len(self.invalid_tokens)} invalid tokens")
        except Exception as e:
            logger.error(f"Error loading whitelist: {e}")
            self.valid_tokens = set()
            self.invalid_tokens = set()
    
    def save_whitelist(self):
        """Save whitelist to file"""
        try:
            os.makedirs("data", exist_ok=True)
            with open(self.whitelist_path, 'w') as f:
                json.dump({
                    "valid_tokens": list(self.valid_tokens),
                    "invalid_tokens": list(self.invalid_tokens),
                    "last_updated": time.time()
                }, f, indent=2)
            logger.info(f"Saved whitelist with {len(self.valid_tokens)} valid tokens and {len(self.invalid_tokens)} invalid tokens")
        except Exception as e:
            logger.error(f"Error saving whitelist: {e}")
    
    async def build_whitelist(self):
        """BYPASS: Skip whitelist building"""
        logger.info("BYPASS: Whitelist building skipped - all tokens will be accepted")
        # Just save the current whitelist
        self.save_whitelist()
        logger.info("BYPASS: Whitelist building complete. All tokens accepted.")
    
    async def run(self):
        """Run the MEV trader"""
        self.running = True
        
        # Setup signal handlers
        for sig in (signal.SIGINT, signal.SIGTERM):
            signal.signal(sig, self._handle_exit)
            
        logger.info("Starting Quantum MEV Trader... ALL VALIDATION BYPASSED - WILL TRADE ANY TOKEN")
        
        # Initial balance check and WSOL account creation
        try:
            # Verify WSOL account exists or create it
            await self.ensure_wsol_account_exists()
            
            # Build token whitelist (bypassed but kept to maintain flow)
            await self.build_whitelist()
            
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            self.last_balance_update = time.time()
            
            self.stats.initial_sol_balance = self.sol_balance + self.wsol_balance
            logger.info(f"Initial wallet balance: {self.stats.initial_sol_balance:.4f} SOL")
        except Exception as e:
            logger.error(f"Error getting initial balance: {e}")
        
        # Start tasks
        tasks = [
            asyncio.create_task(self.monitor_mempool()),
            asyncio.create_task(self.monitor_balances())
        ]
        
        # Keep the bot running
        while self.running:
            await asyncio.sleep(0.1)
            await self.check_active_opportunity()
        
        # Safe shutdown
        await self.safe_shutdown()
        self.print_stats()
        
        # Cancel tasks
        for task in tasks:
            task.cancel()
    
    async def ensure_wsol_account_exists(self):
        """Ensure WSOL token account exists, creating it if needed"""
        wsol_address = os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
        
        # Check if WSOL account is provided in env
        if not self.wsol_account:
            try:
                # Get the default WSOL token account from trader
                self.wsol_account = self.trader.get_associated_token_address(
                    self.wallet_address, wsol_address
                )
                logger.info(f"Using calculated WSOL account: {self.wsol_account}")
            except Exception as e:
                logger.error(f"Error getting WSOL account: {e}")
                return
        
        # Check if the account already exists
        wsol_balance = -1
        try:
            wsol_balance = self.trader.get_token_balance(wsol_address)
        except Exception:
            logger.info("WSOL account does not exist, creating it...")
            
        # If balance check failed, create the WSOL account
        if wsol_balance < 0:
            try:
                # Create the token account
                result = self.trader.create_token_account_if_needed(wsol_address)
                if result:
                    logger.info(f"Successfully created WSOL account: {self.wsol_account}")
                    # Wrap a small amount to initialize the account
                    if self.sol_balance > 0.01:
                        logger.info("Wrapping 0.01 SOL to initialize WSOL account")
                        wrap_result = self.trader.wrap_sol(0.01)
                        if wrap_result:
                            logger.info("Successfully initialized WSOL account")
                            self.last_wrap_time = time.time()
            except Exception as e:
                logger.error(f"Error creating WSOL account: {e}")
    
    def _handle_exit(self, sig, frame):
        """Handle exit signals gracefully"""
        if not self.running:
            return
            
        logger.info("Shutdown signal received, cleaning up...")
        self.running = False
        
    async def safe_shutdown(self):
        """Clean shutdown with completion of active trades"""
        logger.info("Initiating safe shutdown...")
        
        # Complete active opportunity if any
        if self.active_opportunity and self.active_opportunity.status == "frontrun":
            logger.info(f"Completing active opportunity before shutdown")
            await self.execute_emergency_sell(self.active_opportunity)
        
        # Unwrap excess WSOL
        if self.wsol_balance > self.min_wsol_balance * 2:
            await self.unwrap_wsol()
            
        # Save whitelist
        self.save_whitelist()
            
        logger.info("Shutdown complete")
        
    def print_stats(self):
        """Print detailed trading statistics"""
        duration = time.time() - self.stats.start_time
        hours = duration // 3600
        minutes = (duration % 3600) // 60
        seconds = duration % 60
        
        logger.info("=" * 50)
        logger.info("QUANTUM MEV TRADING STATISTICS:")
        logger.info(f"Runtime: {int(hours)}h {int(minutes)}m {int(seconds)}s")
        logger.info(f"Transactions seen: {self.stats.transactions_seen}")
        logger.info(f"Raydium swaps found: {self.stats.raydium_swaps_found}")
        logger.info(f"Viable opportunities: {self.stats.viable_opportunities}")
        logger.info(f"Trades attempted: {self.stats.trades_attempted}")
        logger.info(f"Trades completed: {self.stats.trades_completed}")
        logger.info(f"Failed front-runs: {self.stats.failed_front_runs}")
        logger.info(f"Failed back-runs: {self.stats.failed_back_runs}")
        logger.info(f"Total profit: {self.stats.total_profit:.6f} SOL")
        logger.info(f"Total loss: {self.stats.total_loss:.6f} SOL")
        logger.info(f"Net profit: {self.stats.total_profit - self.stats.total_loss:.6f} SOL")
        
        # Token validation stats
        logger.info(f"Invalid tokens rejected: {self.stats.invalid_tokens_count}")
        logger.info(f"AccountNotFound errors: {self.stats.account_not_found_count}")
        logger.info(f"Whitelist hits: {self.stats.whitelist_hits}")
        
        # GPU stats
        if HAS_TORCH and torch.cuda.is_available():
            logger.info(f"GPU processing time: {self.stats.gpu_processing_time:.2f}s")
            if self.stats.transactions_seen > 0:
                avg_processing = (self.stats.gpu_processing_time * 1000) / self.stats.transactions_seen
                logger.info(f"Average GPU processing time: {avg_processing:.2f}ms/tx")
        
        # Calculate success rate
        if self.stats.trades_attempted > 0:
            success_rate = (self.stats.trades_completed / self.stats.trades_attempted) * 100
            logger.info(f"Success rate: {success_rate:.2f}%")
            
        # Calculate profit per completed trade
        if self.stats.trades_completed > 0:
            profit_per_trade = (self.stats.total_profit - self.stats.total_loss) / self.stats.trades_completed
            logger.info(f"Average profit per completed trade: {profit_per_trade:.6f} SOL")
        
        # WebSocket stats
        logger.info(f"WebSocket connections: {self.stats.ws_connections}")
        logger.info(f"WebSocket messages received: {self.stats.ws_messages_received}")
        logger.info(f"WebSocket errors: {self.stats.ws_errors}")
        logger.info(f"WebSocket reconnects: {self.stats.ws_reconnects}")
        
        logger.info("=" * 50)
    
    async def monitor_mempool(self):
        """Monitor Raydium program logs for MEV opportunities"""
        logger.info(f"Starting mempool monitoring via WebSocket")
        
        retry_count = 0
        backoff = 1.0
        
        while self.running:
            try:
                # Setup WebSocket connection
                headers = {"Content-Type": "application/json"}
                api_key = os.getenv("API_KEY", "")
                if api_key:
                    headers["api-key"] = api_key
                
                # Create a custom SSL context that doesn't verify
                ssl_context = ssl.create_default_context()
                ssl_context.check_hostname = False
                ssl_context.verify_mode = ssl.CERT_NONE
                
                ws_url = os.getenv("WS_URL", "wss://winny-rychu7-fast-mainnet.helius-rpc.com")
                
                # Connect to WebSocket
                async with websockets.connect(
                    ws_url,
                    extra_headers=headers,
                    ping_interval=30,
                    ping_timeout=20,
                    close_timeout=10,
                    max_size=10 * 1024 * 1024,
                    ssl=ssl_context
                ) as ws:
                    self.stats.ws_connections += 1
                    logger.info(f"Connected to WebSocket: {ws_url}")
                    
                    # Reset retry count on successful connection
                    retry_count = 0
                    backoff = 1.0
                    
                    # Subscribe to Raydium program logs
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0", 
                        "id": 1,
                        "method": "logsSubscribe",
                        "params": [
                            {"mentions": [os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8")]},
                            {"commitment": "confirmed"}
                        ]
                    }))
                    
                    # Wait for logs subscription confirmation
                    logs_response = await ws.recv()
                    self.stats.ws_messages_received += 1
                    logs_data = json.loads(logs_response)
                    
                    if "result" not in logs_data:
                        error_msg = logs_data.get("error", {}).get("message", "Unknown error")
                        logger.error(f"Failed to subscribe to logs: {error_msg}")
                        await asyncio.sleep(backoff)
                        backoff = min(30, backoff * 1.5)
                        continue
                    
                    logs_subscription_id = logs_data["result"]
                    logger.info(f"Successfully subscribed to Raydium logs (ID: {logs_subscription_id})")
                    
                    # Subscribe to account updates for balance tracking
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0", 
                        "id": 2,
                        "method": "accountSubscribe",
                        "params": [
                            self.wallet_address,
                            {"commitment": "confirmed", "encoding": "jsonParsed"}
                        ]
                    }))
                    
                    # Wait for account subscription confirmation
                    account_response = await ws.recv()
                    self.stats.ws_messages_received += 1
                    account_data = json.loads(account_response)
                    
                    if "result" in account_data:
                        account_subscription_id = account_data["result"]
                        logger.info(f"Successfully subscribed to wallet account (ID: {account_subscription_id})")
                    else:
                        logger.warning(f"Failed to subscribe to wallet account: {account_data}")
                    
                    # Set up timing variables
                    last_ping_time = time.time()
                    last_heartbeat = time.time()
                    
                    # Process WebSocket messages
                    while self.running:
                        try:
                            # Set shorter timeout to allow for periodic checks
                            message = await asyncio.wait_for(ws.recv(), timeout=5)
                            self.stats.ws_messages_received += 1
                            
                            # Reset heartbeat timer on each message
                            last_heartbeat = time.time()
                            
                            data = json.loads(message)
                            
                            # Check for active opportunity's victim confirmation
                            if self.active_opportunity and self.active_opportunity.status == "frontrun" and not self.active_opportunity.victim_confirmed:
                                if ("method" in data and data["method"] == "logsNotification" and 
                                    "params" in data and "result" in data["params"] and 
                                    "value" in data["params"]["result"] and 
                                    "signature" in data["params"]["result"]["value"]):
                                    
                                    # Check if this is the victim transaction
                                    tx_sig = data["params"]["result"]["value"]["signature"]
                                    if tx_sig == self.active_opportunity.tx_signature:
                                        logger.info(f"VICTIM TRANSACTION CONFIRMED: {tx_sig[:8]}")
                                        self.active_opportunity.victim_confirmed = True
                                        self.stats.victim_confirmed_count += 1
                                        # Execute back-run immediately
                                        await self.execute_back_run(self.active_opportunity)
                                        continue
                            
                            # Handle log notifications (Raydium transactions)
                            if ("method" in data and data["method"] == "logsNotification" and 
                                  "params" in data and "subscription" in data["params"] and 
                                  data["params"]["subscription"] == logs_subscription_id):
                                await self.process_logs_notification(data)
                                
                            # Handle account notifications (balance updates)
                            elif ("method" in data and data["method"] == "accountNotification" and 
                                  "params" in data and "subscription" in data["params"]):
                                await self.update_wallet_balance()
                            
                        except asyncio.TimeoutError:
                            # Check if heartbeat needed
                            current_time = time.time()
                            time_since_last_message = current_time - last_heartbeat
                            
                            # If no messages for twice the heartbeat interval, send ping
                            if time_since_last_message > 20:
                                logger.info(f"No messages for {time_since_last_message:.1f}s, sending heartbeat ping")
                                try:
                                    pong = await ws.ping()
                                    await pong  # Wait for pong response
                                    last_heartbeat = current_time
                                except Exception as e:
                                    logger.warning(f"Heartbeat ping failed: {e}")
                                    # Connection likely dead, break to reconnect
                                    break
                            
                            # Regular checks on timeout
                            await self.check_active_opportunity()
                            
                            # Send regular ping if needed
                            if current_time - last_ping_time > 30:
                                try:
                                    await ws.ping()
                                    last_ping_time = current_time
                                except Exception as e:
                                    logger.warning(f"Regular ping failed: {e}")
                                
                        except websockets.exceptions.ConnectionClosed as e:
                            self.stats.ws_errors += 1
                            logger.warning(f"WebSocket connection closed ({e.code}): {e.reason}, reconnecting...")
                            break
                            
                        except json.JSONDecodeError:
                            self.stats.ws_errors += 1
                            logger.warning("Received invalid JSON, ignoring message")
                            continue
                            
                        except Exception as e:
                            self.stats.ws_errors += 1
                            logger.error(f"Error processing message: {e}")
                            logger.error(traceback.format_exc())
                
                # Connection closed, now reconnect
                logger.info("WebSocket connection ended, reconnecting...")
                self.stats.ws_reconnects += 1
            
            except asyncio.TimeoutError:
                self.stats.ws_errors += 1
                logger.warning(f"WebSocket connection timeout, reconnecting...")
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            except websockets.exceptions.WebSocketException as e:
                self.stats.ws_errors += 1
                retry_count += 1
                logger.error(f"WebSocket error (retry {retry_count}): {e}")
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            except Exception as e:
                self.stats.ws_errors += 1
                retry_count += 1
                logger.error(f"Unexpected error: {e}")
                logger.error(traceback.format_exc())
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            # Always try to reconnect regardless of error
            if self.running:
                logger.info("Attempting to reconnect WebSocket...")
    
    async def update_wallet_balance(self):
        """Update wallet balance after transaction"""
        async with self.balance_lock:
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            self.last_balance_update = time.time()
            
            logger.info(f"Updated wallet balance: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
            
            # Check for victim confirmation if we have active opportunity
            if (self.active_opportunity and 
                self.active_opportunity.status == "frontrun" and 
                not self.active_opportunity.victim_confirmed):
                
                # Check if balance updates suggest victim transaction went through
                # This is a backup method in case we don't see the exact tx in mempool
                logger.info(f"Balance update detected - checking if this could be victim confirmation")
                await asyncio.sleep(0.1)  # Small delay
                self.active_opportunity.victim_confirmed = True
                await self.execute_back_run(self.active_opportunity)
    
    async def process_logs_notification(self, data):
        """Process Raydium program logs for MEV opportunities"""
        try:
            # Skip if shutting down or already have active opportunity
            if not self.running or self.active_opportunity:
                return
                
            # Extract logs and signature
            params = data.get("params", {})
            result = params.get("result", {})
            value = result.get("value", {})
            
            logs = value.get("logs", [])
            tx_signature = value.get("signature", "unknown")
            
            # Skip if no logs or already processed
            if not logs or tx_signature in self.processed_txs:
                return
                
            # Start GPU processing time measurement
            gpu_start_time = time.time()
            
            self.stats.transactions_seen += 1
            self.processed_txs.add(tx_signature)
            
            # Limit processed_txs size
            if len(self.processed_txs) > 5000:
                self.processed_txs = set(list(self.processed_txs)[-2500:])
            
            # Check for ray_log entries with GPU acceleration
            ray_log_found = False
            token_mint = None
            accounts = {}
            amount_sol = 0.0
            is_buy = False
            
            for log in logs:
                if "ray_log:" in log:
                    ray_log_found = True
                    ray_data = self.token_extractor.decode_ray_log(log)
                    
                    if ray_data and ray_data["instruction_type"] in [3, 9]:  # Both 0x03 and 0x09 indicate buys
                        self.stats.raydium_swaps_found += 1
                        
                        # Check amount to see if it's a viable opportunity
                        amount_sol = ray_data["amount"]
                        is_buy = ray_data["is_buy"]
                        
                        if (is_buy and 
                            self.min_victim_size <= amount_sol <= self.max_victim_size):
                            
                            # Extract token mint using GPU acceleration
                            if HAS_TORCH and torch.cuda.is_available() and hasattr(self.token_extractor, 'device') and self.token_extractor.device:
                                # GPU-accelerated token extraction
                                token_mint = self.token_extractor.accelerate_log_parsing(logs)
                            else:
                                # CPU fallback
                                token_mint = await asyncio.get_event_loop().run_in_executor(
                                    self.thread_pool,
                                    lambda: self.token_extractor.extract_token_mint(logs)
                                )
                            
                            if token_mint:
                                # Extract accounts for swap
                                accounts = self.token_extractor.extract_accounts(logs)
                                accounts["token_mint"] = token_mint
                                
                                logger.info(f"OPPORTUNITY FOUND: {tx_signature[:8]}, "
                                          f"Buy {amount_sol:.4f} SOL, "
                                          f"Token: {token_mint[:8]}")
                                          
                                self.stats.viable_opportunities += 1
                                          
                                # Update GPU time stats
                                gpu_processing_time = time.time() - gpu_start_time
                                self.stats.gpu_processing_time += gpu_processing_time
                                
                                # Process opportunity - BYPASS VALIDATION
                                await self.process_opportunity(token_mint, amount_sol, tx_signature, accounts)
                        
                        else:
                            # Log if it's a Raydium swap but not a viable opportunity
                            token_info = "Unknown"
                            if ray_data["pubkeys"] and len(ray_data["pubkeys"]) > 3:
                                possible_token = ray_data["pubkeys"][3]
                                if self.token_extractor.is_valid_pubkey(possible_token):
                                    token_info = possible_token[:8]
                                    
                            logger.debug(f"Raydium swap (not viable): {tx_signature[:8]}, "
                                    f"Buy: {is_buy}, Amount: {amount_sol:.4f} SOL, "
                                    f"Token: {token_info}")
            
            if not ray_log_found:
                logger.debug(f"Transaction without ray_log: {tx_signature[:8]}")
                
        except Exception as e:
            logger.error(f"Error processing logs: {e}")
            logger.error(traceback.format_exc())
    
    async def check_active_opportunity(self):
        """Check and update active opportunity status"""
        if not self.active_opportunity:
            return
            
        # Handle timeout for victim confirmation
        timeout = int(os.getenv("VICTIM_CONFIRM_TIMEOUT", "15"))
        if (self.active_opportunity.status == "frontrun" and 
            not self.active_opportunity.victim_confirmed and
            time.time() - self.active_opportunity.start_time > timeout):
            
            logger.warning(f"Victim tx confirmation timeout for {self.active_opportunity.token_mint[:8]}")
            self.stats.victim_missed_count += 1
            
            # Execute emergency sell
            await self.execute_emergency_sell(self.active_opportunity)
    
    async def process_opportunity(self, token_mint, amount_sol, tx_signature, accounts):
        """MODIFIED: Process an MEV opportunity without validation"""
        try:
            # Skip if already have active opportunity or shutting down
            if self.active_opportunity or not self.running:
                return
            
            # Calculate front-run amount based on victim size but with safety limit 
            front_amount = min(self.front_run_amount, self.sol_balance * 0.2)
            
            # Create opportunity object
            opportunity = TradeOpportunity(
                token_mint=token_mint,
                amount_sol=amount_sol,
                tx_signature=tx_signature,
                accounts=accounts,
                front_amount=front_amount,
                start_time=time.time(),
                start_sol_balance=self.sol_balance,
                start_wsol_balance=self.wsol_balance
            )
            
            # Check if we have enough balance
            async with self.balance_lock:
                total_balance = self.sol_balance + self.wsol_balance
                min_required = front_amount + self.min_sol_balance
                
                if total_balance < min_required:
                    logger.warning(f"Insufficient balance: {total_balance:.4f} SOL, need {min_required:.4f} SOL")
                    return
                    
                # Make sure we have enough WSOL
                wsol_ok = await self.ensure_wsol(front_amount * 1.2)
                
                if not wsol_ok:
                    logger.error(f"Failed to ensure WSOL for {token_mint[:8]}")
                    return
                    
                # Update WSOL balance after wrapping
                self.wsol_balance = self.trader.get_token_balance(
                    os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                )
            
            # Set as active opportunity
            self.active_opportunity = opportunity
            self.stats.trades_attempted += 1
            
            logger.info(f"EXECUTING FRONT-RUN: Victim tx {tx_signature[:8]} with {front_amount:.4f} SOL -> {token_mint[:8]}")
            
            # Execute front-run
            front_tx_signature = self.execute_swap(
                token_mint=token_mint,
                amount_sol=front_amount,
                accounts=accounts,
                is_buy=True
            )
            
            if front_tx_signature:
                logger.info(f"Front-run successful: {front_tx_signature[:8]}")
                opportunity.front_tx = front_tx_signature
                opportunity.status = "frontrun"
                
                # Get token balance and decimals for accurate back-run
                opportunity.token_balance = self.trader.get_token_balance(token_mint)
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
                
                logger.info(f"Received {opportunity.token_balance} tokens (decimals: {opportunity.token_decimals})")
                
                # Now wait for victim confirmation via WebSocket balance updates
            else:
                logger.error(f"Front-run failed: null tx")
                self.stats.failed_front_runs += 1
                self.active_opportunity = None
                
        except Exception as e:
            logger.error(f"Error processing opportunity: {e}")
            logger.error(traceback.format_exc())
            self.active_opportunity = None
    
    def execute_swap(self, token_mint, amount_sol, accounts, is_buy):
        """Execute a swap transaction with correct instruction format"""
        try:
            # Create instruction data with proper 0x09 format
            instr_data = bytearray(17)  # 17 bytes: 1 byte discriminator + 8 bytes amount + 8 bytes min_out
            
            # Set discriminator to 0x09 (Raydium V4 swap)
            instr_data[0] = 0x09
            
            # Calculate amount in lamports
            amount_lamports = int(amount_sol * LAMPORTS_PER_SOL)
            
            # Convert amount to little-endian bytes
            amount_bytes = amount_lamports.to_bytes(8, byteorder="little")
            
            # Calculate min_out with slippage
            slippage = float(os.getenv("SLIPPAGE", "3.0")) / 100.0
            min_out_lamports = int(amount_lamports * (1.0 - slippage))
            min_out_bytes = min_out_lamports.to_bytes(8, byteorder="little")
            
            # Set amount_in (bytes 1-8)
            for i in range(8):
                instr_data[i+1] = amount_bytes[i]
                
            # Set min_out (bytes 9-16)
            for i in range(8):
                instr_data[i+9] = min_out_bytes[i]
            
            # Execute swap with Rust trading core
            tx_signature = self.trader.swap_with_accounts(
                instr_data,  # Proper 0x09 instruction data
                accounts.get("amm_id", ""),
                accounts.get("amm_authority", ""),
                accounts.get("amm_open_orders", ""),
                accounts.get("amm_target_orders", ""),
                accounts.get("pool_coin_token_account", ""),
                accounts.get("pool_pc_token_account", ""),
                accounts.get("serum_market", ""),
                accounts.get("serum_bids", ""),
                accounts.get("serum_asks", ""),
                accounts.get("serum_event_queue", ""),
                accounts.get("serum_coin_vault", ""),
                accounts.get("serum_pc_vault", ""),
                accounts.get("serum_vault_signer", ""),
                token_mint,
                is_buy
            )
            
            return tx_signature
        
        except Exception as e:
            logger.error(f"Error executing swap: {e}")
            return None
    
    def execute_token_sell(self, token_mint, token_amount, token_decimals, accounts):
        """Sell tokens with proper instruction format"""
        try:
            # Create instruction data with proper 0x09 format
            instr_data = bytearray(17)  # 17 bytes: 1 byte discriminator + 8 bytes amount + 8 bytes min_out
            
            # Set discriminator to 0x09 (Raydium V4 swap)
            instr_data[0] = 0x09
            
            # Calculate token amount in smallest units
            token_units = int(token_amount * (10 ** token_decimals))
            
            # Convert amount to little-endian bytes
            amount_bytes = token_units.to_bytes(8, byteorder="little")
            
            # Calculate min_out with slippage (higher for sell)
            slippage = float(os.getenv("SLIPPAGE", "3.0")) / 100.0  # Use same slippage for sell
            min_out_lamports = 0  # For selling, we set min_out to 0 for speed
            min_out_bytes = min_out_lamports.to_bytes(8, byteorder="little")
            
            # Set amount_in (bytes 1-8)
            for i in range(8):
                instr_data[i+1] = amount_bytes[i]
                
            # Set min_out (bytes 9-16)
            for i in range(8):
                instr_data[i+9] = min_out_bytes[i]
            
            # Execute sell with Rust trading core
            tx_signature = self.trader.swap_with_accounts(
                instr_data,  # Proper 0x09 instruction data
                accounts.get("amm_id", ""),
                accounts.get("amm_authority", ""),
                accounts.get("amm_open_orders", ""),
                accounts.get("amm_target_orders", ""),
                accounts.get("pool_coin_token_account", ""),
                accounts.get("pool_pc_token_account", ""),
                accounts.get("serum_market", ""),
                accounts.get("serum_bids", ""),
                accounts.get("serum_asks", ""),
                accounts.get("serum_event_queue", ""),
                accounts.get("serum_coin_vault", ""),
                accounts.get("serum_pc_vault", ""),
                accounts.get("serum_vault_signer", ""),
                token_mint,
                False  # Is sell
            )
            
            return tx_signature
        
        except Exception as e:
            logger.error(f"Error selling tokens: {e}")
            return None
    
    async def execute_back_run(self, opportunity):
        """Execute back-run after victim tx confirmation"""
        try:
            token_mint = opportunity.token_mint
            
            # Get token balance with verification
            if opportunity.token_balance <= 0:
                opportunity.token_balance = self.trader.get_token_balance(token_mint)
            
            if opportunity.token_balance <= 0:
                logger.warning(f"No tokens to sell for back-run: {token_mint[:8]}")
                self.active_opportunity = None
                return
                
            logger.info(f"EXECUTING BACK-RUN: Selling {opportunity.token_balance} tokens of {token_mint[:8]}")
            
            # Get token decimals if not already known
            if opportunity.token_decimals == 9:  # Default value
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
            
            # Execute back-run immediately - no sleep delay needed as victim tx is confirmed
            back_tx_signature = self.execute_token_sell(
                token_mint=token_mint,
                token_amount=opportunity.token_balance,
                token_decimals=opportunity.token_decimals,
                accounts=opportunity.accounts
            )
            
            if back_tx_signature:
            
            logger.info(f"Back-run successful: {back_tx_signature[:8]}")
                opportunity.back_tx = back_tx_signature
    
    # Get updated balances
    async with self.balance_lock:
        self.sol_balance = self.trader.get_sol_balance()
        self.wsol_balance = self.trader.get_token_balance(
            os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
        )
        self.last_balance_update = time.time()
        
        # Calculate profit/loss
        total_end = self.sol_balance + self.wsol_balance
        total_start = opportunity.start_sol_balance + opportunity.start_wsol_balance
        
        # Calculate transaction costs
        tx_cost = 0.000005 * 2  # Base fee for two transactions
        compute_cost = (int(os.getenv("COMPUTE_LIMIT", "50000")) / 1000000) * (int(os.getenv("PRIORITY_FEE", "750000")) / LAMPORTS_PER_SOL) * 2
        total_cost = tx_cost + compute_cost
        
        raw_profit = total_end - total_start
        net_profit = raw_profit - total_cost
        
        opportunity.profit = net_profit
        opportunity.end_sol_balance = self.sol_balance
        opportunity.end_wsol_balance = self.wsol_balance
        opportunity.status = "complete"
        
        self.stats.trades_completed += 1
        
        if net_profit >= 0:
            self.stats.total_profit += net_profit
            roi_percent = (net_profit / opportunity.front_amount) * 100
            logger.info(f"TRADE PROFIT: +{net_profit:.6f} SOL (ROI: {roi_percent:.2f}%)")
        else:
            self.stats.total_loss += abs(net_profit)
            logger.info(f"TRADE LOSS: {net_profit:.6f} SOL")
else:
    logger.error(f"Back-run failed: null tx")
    self.stats.failed_back_runs += 1
    await self.execute_emergency_sell(opportunity)
    return

# Clear active opportunity
self.active_opportunity = None
    
except Exception as e:
    logger.error(f"Error in back-run: {e}")
    logger.error(traceback.format_exc())
    await self.execute_emergency_sell(opportunity)

async def execute_emergency_sell(self, opportunity):
    """Emergency sell to recover funds"""
    try:
        token_mint = opportunity.token_mint
        
        # Get current token balance
        token_balance = self.trader.get_token_balance(token_mint)
        
        if token_balance <= 0:
            logger.warning(f"No tokens to sell: {token_mint[:8]}")
            self.active_opportunity = None
            return
            
        logger.info(f"EMERGENCY SELLING {token_balance} tokens of {token_mint[:8]}")
        
        # Get token decimals if not already known
        if opportunity.token_decimals == 9:  # Default value
            try:
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
            except Exception:
                # If can't get decimals, assume 9
                opportunity.token_decimals = 9
        
        # Execute sell with higher slippage
        for attempt in range(3):
            multiplier = 1.0 + (attempt * 0.5)  # Increase slippage with each attempt
            
            tx_signature = self.execute_token_sell(
                token_mint=token_mint,
                token_amount=token_balance,
                token_decimals=opportunity.token_decimals,
                accounts=opportunity.accounts
            )
            
            if tx_signature:
                logger.info(f"Emergency sell successful on attempt {attempt+1}: {tx_signature[:8]}")
                
                # Get updated balances
                async with self.balance_lock:
                    self.sol_balance = self.trader.get_sol_balance()
                    self.wsol_balance = self.trader.get_token_balance(
                        os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                    )
                    self.last_balance_update = time.time()
                    
                    # Calculate profit/loss
                    total_end = self.sol_balance + self.wsol_balance
                    profit = total_end - (opportunity.start_sol_balance + opportunity.start_wsol_balance)
                    
                    opportunity.profit = profit
                    opportunity.end_sol_balance = self.sol_balance
                    opportunity.end_wsol_balance = self.wsol_balance
                    
                    if profit >= 0:
                        self.stats.total_profit += profit
                        logger.info(f"Emergency sell profit: +{profit:.6f} SOL")
                    else:
                        self.stats.total_loss += abs(profit)
                        logger.info(f"Emergency sell loss: {profit:.6f} SOL")
                
                break
            else:
                logger.error(f"Emergency sell attempt {attempt+1} failed")
                await asyncio.sleep(0.5 * (attempt + 1))
        
        # Clear active opportunity
        self.active_opportunity = None
            
    except Exception as e:
        logger.error(f"Error in emergency sell: {e}")
        logger.error(traceback.format_exc())
        self.active_opportunity = None

def estimate_profit(self, front_amount, victim_amount):
    """Estimate profit with enhanced calculation using GPU acceleration if available"""
    try:
        if HAS_TORCH and torch.cuda.is_available() and hasattr(self, 'device') and self.device:
            # Use GPU for more complex profit calculation
            with torch.cuda.stream(self.token_extractor.stream if hasattr(self.token_extractor, 'stream') else None):
                # Raydium fee
                raydium_fee = 0.0025  # 0.25%
                
                # Create tensors for calculation
                front_tensor = torch.tensor([front_amount], device=self.device, dtype=torch.float32)
                victim_tensor = torch.tensor([victim_amount], device=self.device, dtype=torch.float32)
                
                # Calculate token flow
                fee_factor = torch.tensor([1.0 - raydium_fee], device=self.device)
                tokens_received = front_tensor * fee_factor
                
                # Calculate price impact based on victim size
                # Complex model: price impact increases with square root of victim size
                # but is capped at 20%
                base_impact = torch.tensor([0.05], device=self.device)
                normalized_size = victim_tensor / 5.0
                sqrt_size = torch.sqrt(normalized_size)
                price_impact = torch.min(base_impact * sqrt_size, torch.tensor([0.20], device=self.device))
                
                # Profit calculation
                impact_factor = torch.tensor([1.0], device=self.device) + price_impact
                back_amount = tokens_received * impact_factor * fee_factor
                
                # Transaction fees
                priority_fee = torch.tensor([(int(os.getenv("PRIORITY_FEE", "750000")) / LAMPORTS_PER_SOL) * 
                                             (int(os.getenv("COMPUTE_LIMIT", "50000")) / 1000000) * 2], device=self.device)
                base_fee = torch.tensor([0.000005 * 2], device=self.device)  # Two transactions
                total_fees = priority_fee + base_fee
                
                # Final profit calculations
                raw_profit = back_amount - front_tensor
                net_profit = raw_profit - total_fees
                adjusted_profit = net_profit * 0.95  # Safety margin
                
                # Convert results back to CPU
                torch.cuda.synchronize()
                result = {
                    "raw_profit": raw_profit.item(),
                    "net_profit": net_profit.item(),
                    "adjusted_profit": adjusted_profit.item(),
                    "total_fees": total_fees.item(),
                    "price_impact": price_impact.item(),
                    "is_profitable": adjusted_profit.item() > self.min_profit
                }
                
                return result
        
        # CPU fallback if GPU not available
        # Raydium fee
        raydium_fee = 0.0025  # 0.25%
        
        # Calculate token flow
        tokens_received = front_amount * (1.0 - raydium_fee)
        
        # Calculate price impact based on victim size
        # Larger victims cause more price impact
        price_impact_percent = min(0.05 * (victim_amount / 5.0) ** 0.5, 0.20)
        
        # Calculate expected SOL from back-run
        back_amount = tokens_received * (1.0 + price_impact_percent) * (1.0 - raydium_fee)
        
        # Calculate total fees
        priority_fee = (int(os.getenv("PRIORITY_FEE", "750000")) / LAMPORTS_PER_SOL) * (int(os.getenv("COMPUTE_LIMIT", "50000")) / 1000000) * 2
        total_fees = priority_fee + (0.000005 * 2)  # Base fee for two transactions
        
        # Calculate profit after fees
        raw_profit = back_amount - front_amount
        net_profit = raw_profit - total_fees
        
        # Apply safety margin
        adjusted_profit = net_profit * 0.95
        
        return {
            "raw_profit": raw_profit,
            "net_profit": net_profit,
            "adjusted_profit": adjusted_profit,
            "total_fees": total_fees,
            "price_impact": price_impact_percent,
            "is_profitable": adjusted_profit > self.min_profit
        }
    except Exception as e:
        logger.error(f"Error in profit estimation: {e}, falling back to conservative estimate")
        # Ultra-conservative fallback
        return {
            "raw_profit": 0,
            "net_profit": 0,
            "adjusted_profit": 0,
            "total_fees": 0.001,
            "price_impact": 0,
            "is_profitable": False
        }

async def ensure_wsol(self, amount_needed):
    """Ensure sufficient WSOL for trading with reduced requirements"""
    # Already have enough
    if self.wsol_balance >= amount_needed:
        return True
        
    # Don't wrap too frequently
    if time.time() - self.last_wrap_time < self.wsol_wrap_interval:
        return False
        
    # Calculate amount to wrap (reduced buffer)
    wrap_amount = max(
        0.01,  # Minimum wrap amount reduced to 0.01
        (amount_needed - self.wsol_balance) * 1.1  # Buffer reduced from 1.5 to 1.1
    )
    
    # Ensure we have enough SOL - reduced min balance to improve success rate
    if self.sol_balance < wrap_amount + (self.min_sol_balance / 2):  # Reduced min balance requirement
        logger.error(f"Insufficient SOL: {self.sol_balance}, need {wrap_amount + (self.min_sol_balance / 2)}")
        return False
        
    try:
        logger.info(f"Wrapping {wrap_amount} SOL to WSOL")
        
        # Execute wrap
        signature = self.trader.wrap_sol(wrap_amount)
        self.last_wrap_time = time.time()
        
        # Wait for confirmation
        confirmed = self.trader.confirm_transaction(signature, 15)
        
        if confirmed:
            # Update balances
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            logger.info(f"Wrapped SOL successfully: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
            return True
        else:
            logger.error(f"Wrap confirmation failed")
            return False
            
    except Exception as e:
        logger.error(f"Error wrapping SOL: {e}")
        return False

async def unwrap_wsol(self):
    """Unwrap WSOL back to native SOL"""
    try:
        logger.info(f"Unwrapping {self.wsol_balance} WSOL to SOL")
        
        # Execute unwrap
        signature = self.trader.unwrap_sol()
        
        # Wait for confirmation
        confirmed = self.trader.confirm_transaction(signature, 15)
        
        if confirmed:
            # Update balances
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            logger.info(f"Unwrapped WSOL successfully: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
            return True
        else:
            logger.error(f"Unwrap confirmation failed")
            return False
            
    except Exception as e:
        logger.error(f"Error unwrapping WSOL: {e}")
        return False

async def monitor_balances(self):
    """Monitor and manage SOL/WSOL balances"""
    check_interval = 60  # Check every minute
    
    while self.running:
        try:
            # Skip if active opportunity
            if self.active_opportunity:
                await asyncio.sleep(check_interval)
                continue
            
            # Get current balances
            current_time = time.time()
            if current_time - self.last_balance_update > 60:
                self.sol_balance = self.trader.get_sol_balance()
                self.wsol_balance = self.trader.get_token_balance(
                    os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                )
                self.last_balance_update = current_time
            
            # Check if we need to wrap SOL or unwrap WSOL
            total_balance = self.sol_balance + self.wsol_balance
            wsol_ratio = self.wsol_balance / total_balance if total_balance > 0 else 0
            target_ratio = float(os.getenv("WSOL_TARGET_RATIO", "0.75"))
            
            # If SOL balance is too low, unwrap WSOL
            if self.sol_balance < self.min_sol_balance and self.wsol_balance > self.min_sol_balance * 2:
                logger.info(f"SOL balance too low ({self.sol_balance:.4f}), unwrapping WSOL")
                await self.unwrap_wsol()
            
            # If WSOL ratio is too low, wrap SOL
            elif wsol_ratio < target_ratio * 0.8 and self.sol_balance > self.min_sol_balance * 1.5:
                wrap_amount = (total_balance * target_ratio - self.wsol_balance) * 1.1
                if wrap_amount > 0.01:  # Reduced min wrap amount
                    logger.info(f"WSOL ratio too low ({wsol_ratio:.2f}), wrapping {wrap_amount:.4f} SOL")
                    await self.ensure_wsol(wrap_amount)
            
        except Exception as e:
            logger.error(f"Error in balance monitoring: {e}")
            
        await asyncio.sleep(check_interval)

async def main():
    """Main entry point"""
    try:
        # Initialize and run the MEV trader
        trader = QuantumMEVTrader()
        await trader.run()
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received, shutting down gracefully...")
    except Exception as e:
        logger.critical(f"Critical error: {e}")
        logger.critical(traceback.format_exc())

if __name__ == "__main__":
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())