import asyncio
import base64
import json
import logging
import time
import websockets
from typing import Dict, List, Optional, Set, Tuple
from decimal import Decimal
import torch
from dataclasses import dataclass, field
from collections import deque

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

@dataclass
class TradingMetrics:
    messages_received: int = 0
    opportunities_found: int = 0
    successful_trades: int = 0
    total_profit: Decimal = Decimal('0')
    start_time: float = field(default_factory=time.time)
    sandwich_profits: Decimal = Decimal('0')
    arbitrage_profits: Decimal = Decimal('0')
    failed_attempts: Dict[str, int] = field(default_factory=dict)
    active_paths: Set[str] = field(default_factory=set)

class AdvancedTrading:
    def __init__(self, trader, config):
        self.trader = trader
        self.config = config
        self.active = False
        self.metrics = TradingMetrics()
        
        # Initialize CUDA streams for parallel processing
        self.device = torch.device('cuda')
        self.streams = {
            'arbitrage': torch.cuda.Stream(),
            'sandwich': torch.cuda.Stream(),
            'direct': torch.cuda.Stream(),
            'compute': torch.cuda.Stream()
        }
        
        # Setup processing buffers
        self.setup_processing_buffers()
        logger.info("Advanced trading initialized with GPU acceleration")

    def setup_processing_buffers(self):
        """Initialize GPU processing buffers"""
        try:
            with torch.cuda.stream(self.streams['compute']):
                self.opportunity_buffer = torch.zeros(
                    (1024, 128), 
                    device=self.device, 
                    dtype=torch.float16
                ).pin_memory()
                
                self.price_buffer = torch.zeros(
                    (512, 2), 
                    device=self.device, 
                    dtype=torch.float32
                ).pin_memory()
                
                logger.debug("GPU buffers initialized")
        except Exception as e:
            logger.error(f"Buffer initialization error: {e}")
            raise

    async def start_websocket_monitoring(self):
        """Start monitoring WebSocket for opportunities"""
        self.active = True
        logger.info("Starting advanced trading monitor...")
        
        while self.active:
            try:
                async with websockets.connect(
                    self.config['rpc']['helius']['websocket'],
                    ping_interval=self.config['system']['websocket']['ping_interval'],
                    ping_timeout=self.config['system']['websocket']['ping_timeout'],
                    close_timeout=self.config['system']['websocket']['close_timeout'],
                    max_size=self.config['system']['websocket']['max_message_size']
                ) as ws:
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "logsSubscribe",
                        "params": [
                            {"mentions": [self.config['programs']['raydium']]},
                            {"commitment": "processed"}
                        ]
                    }))
                    
                    while self.active:
                        msg = await ws.recv()
                        self.metrics.messages_received += 1
                        
                        # Process opportunities in parallel
                        await self._process_message(msg)
                        
                        # Log metrics periodically
                        if self.metrics.messages_received % 100 == 0:
                            self._log_metrics()
                            
            except websockets.exceptions.ConnectionClosed:
                logger.warning("WebSocket connection closed. Reconnecting...")
                await asyncio.sleep(1)
            except Exception as e:
                logger.error(f"WebSocket error: {e}")
                await asyncio.sleep(1)

    async def _process_message(self, msg: str):
        """Process incoming WebSocket message with parallel strategy checking"""
        try:
            data = json.loads(msg)
            if 'params' not in data:
                return

            logs = data['params'].get('result', {}).get('value', {}).get('logs', [])
            for log in logs:
                if "ray_log:" not in log:
                    continue

                raw_data = base64.b64decode(log.split('ray_log: ')[1].strip())
                self.metrics.opportunities_found += 1

                # Check all strategies in parallel with GPU acceleration
                results = await asyncio.gather(
                    self._check_sandwich_opportunity(raw_data),
                    self._check_triangle_arbitrage(raw_data),
                    self._check_multi_hop_opportunity(raw_data),
                    self._check_direct_swap(raw_data),
                    return_exceptions=True
                )

                valid_opportunities = [r for r in results if isinstance(r, dict) and r.get('profit', 0) > 0]
                
                if valid_opportunities:
                    # Sort by profit/risk ratio
                    best_opp = max(valid_opportunities, key=lambda x: 
                        x['profit'] * x.get('confidence', 1.0) / x.get('impact', 1.0)
                    )
                    
                    logger.info(f"Found opportunity: {best_opp['type']} with profit {best_opp['profit']:.4f}%")
                    await self._execute_opportunity(best_opp)

        except Exception as e:
            logger.error(f"Message processing error: {e}")

    async def _check_sandwich_opportunity(self, raw_data: bytes) -> Optional[Dict]:
        """Check for sandwich opportunities with MEV protection"""
        try:
            with torch.cuda.stream(self.streams['sandwich']):
                victim_tx = await self._decode_victim_transaction(raw_data)
                if not victim_tx:
                    return None

                pool_address = victim_tx['pool']
                
                # Calculate optimal positions
                front_run = await self._calculate_optimal_front_run(
                    pool_address,
                    victim_tx['amount'],
                    victim_tx['direction']
                )

                if not front_run:
                    return None

                back_run = await self._calculate_optimal_back_run(
                    pool_address,
                    front_run['amount'],
                    victim_tx['amount']
                )

                if not back_run:
                    return None

                # Verify MEV protection
                mev_safe = await self.trader.mev_protection.verify_sandwich(
                    front_run['amount'],
                    victim_tx['amount'],
                    back_run['amount'],
                    pool_address
                )

                if not mev_safe:
                    return None

                # Calculate total profit
                profit = await self._calculate_sandwich_profit(
                    front_run,
                    back_run,
                    pool_address
                )

                return {
                    'type': 'sandwich',
                    'profit': profit,
                    'front_run': front_run,
                    'back_run': back_run,
                    'victim': victim_tx,
                    'pool': pool_address,
                    'confidence': 0.9  # High confidence for sandwich
                }

        except Exception as e:
            logger.error(f"Sandwich check error: {e}")
            return None

    async def _check_triangle_arbitrage(self, raw_data: bytes) -> Optional[Dict]:
        """Check for triangle arbitrage opportunities"""
        try:
            with torch.cuda.stream(self.streams['arbitrage']):
                # Get token info
                token_info = await self._decode_token_info(raw_data)
                if not token_info:
                    return None

                # Find optimal triangle path
                optimal_path = await self._find_optimal_triangle(
                    token_info['token_in'],
                    self.config['trading']['position']['size']
                )

                if not optimal_path:
                    return None

                # Verify execution is possible
                execution_check = await self._verify_triangle_execution(
                    optimal_path['path'],
                    optimal_path['amounts']
                )

                if not execution_check['executable']:
                    return None

                return {
                    'type': 'triangle',
                    'profit': optimal_path['profit'],
                    'path': optimal_path['path'],
                    'amounts': optimal_path['amounts'],
                    'impact': optimal_path['impact'],
                    'confidence': execution_check['confidence']
                }

        except Exception as e:
            logger.error(f"Triangle arbitrage error: {e}")
            return None

    async def _check_multi_hop_opportunity(self, raw_data: bytes) -> Optional[Dict]:
        """Check for multi-hop opportunities with optimal routing"""
        try:
            # Decode token information
            token_info = await self._decode_token_info(raw_data)
            if not token_info:
                return None

            # Find optimal route
            optimal_route = await self.trader.impact_calculator.optimize_route(
                self.config['trading']['position']['size'],
                token_info['token_in'],
                token_info['token_out']
            )

            if not optimal_route:
                return None

            # Validate profitability
            if optimal_route['impact'] > self.config['trading']['limits']['max_price_impact']:
                return None

            return {
                'type': 'multi_hop',
                'profit': optimal_route['profit'],
                'path': optimal_route['path'],
                'amounts': optimal_route['amounts'],
                'impact': optimal_route['impact'],
                'confidence': 0.7
            }

        except Exception as e:
            logger.error(f"Multi-hop check error: {e}")
            return None

    async def _check_direct_swap(self, raw_data: bytes) -> Optional[Dict]:
        """Check for direct swap opportunities"""
        try:
            with torch.cuda.stream(self.streams['direct']):
                opportunity = await self.trader.gpu_accelerator.analyze_opportunity(raw_data)
                if not opportunity:
                    return None

                # Get pattern matching
                similar_patterns = await self.trader.vector_store.search_similar_pools(
                    opportunity,
                    k=3,
                    min_similarity=0.7
                )

                # Calculate confidence based on patterns
                confidence = self._calculate_confidence(opportunity, similar_patterns)

                if confidence < 0.6:
                    return None

                return {
                    'type': 'direct',
                    'profit': opportunity['profit'],
                    'amount_in': opportunity['amount_in'],
                    'amount_out': opportunity['amount_out'],
                    'impact': opportunity['impact'],
                    'confidence': confidence
                }

        except Exception as e:
            logger.error(f"Direct swap check error: {e}")
            return None

    async def _execute_opportunity(self, opportunity: Dict) -> bool:
        """Execute trading opportunity with type-specific handling"""
        try:
            logger.info(f"Executing {opportunity['type']} opportunity...")
            result = False

            if opportunity['type'] == 'sandwich':
                result = await self._execute_sandwich_trade(opportunity)
            elif opportunity['type'] == 'triangle':
                result = await self._execute_triangle_trade(opportunity)
            elif opportunity['type'] == 'multi_hop':
                result = await self._execute_multi_hop_trade(opportunity)
            else:  # direct swap
                result = await self._execute_direct_trade(opportunity)

            if result:
                self.metrics.successful_trades += 1
                profit = Decimal(str(opportunity['profit']))
                self.metrics.total_profit += profit

                if opportunity['type'] == 'sandwich':
                    self.metrics.sandwich_profits += profit
                elif opportunity['type'] in ['triangle', 'multi_hop']:
                    self.metrics.arbitrage_profits += profit

                logger.info(f"""
Trade executed successfully:
Type: {opportunity['type']}
Profit: {opportunity['profit']:.4f}%
Total Profit: {float(self.metrics.total_profit):.6f} SOL
Total Trades: {self.metrics.successful_trades}
                """)

            return result

        except Exception as e:
            logger.error(f"Trade execution error: {e}")
            return False

    async def _execute_sandwich_trade(self, opportunity: Dict) -> bool:
        """Execute sandwich trade with precise timing"""
        try:
            # Execute front-run
            front_run_result = await self.trader.trading_core.execute_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['front_run']['amount'],
                min_amount_out=opportunity['front_run']['min_out']
            )

            if not front_run_result:
                return False

            # Wait for victim transaction
            await asyncio.sleep(0.1)  # Minimal delay

            # Execute back-run
            back_run_result = await self.trader.trading_core.execute_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['back_run']['amount'],
                min_amount_out=opportunity['back_run']['min_out']
            )

            return bool(back_run_result)

        except Exception as e:
            logger.error(f"Sandwich execution error: {e}")
            return False

    async def _execute_triangle_trade(self, opportunity: Dict) -> bool:
        """Execute triangle arbitrage trade"""
        try:
            for i, pool in enumerate(opportunity['path']):
                result = await self.trader.trading_core.execute_trade(
                    pool_address=pool,
                    amount_in=opportunity['amounts'][i],
                    min_amount_out=opportunity['amounts'][i+1] * (1 - self.config['execution']['slippage'])
                )
                
                if not result:
                    return False
                    
                await asyncio.sleep(0.1)  # Small delay between hops
                
            return True

        except Exception as e:
            logger.error(f"Triangle trade execution error: {e}")
            return False

    async def _execute_multi_hop_trade(self, opportunity: Dict) -> bool:
        """Execute multi-hop trade"""
        try:
            for i, pool in enumerate(opportunity['path']):
                result = await self.trader.trading_core.execute_trade(
                    pool_address=pool,
                    amount_in=opportunity['amounts'][i],
                    min_amount_out=opportunity['amounts'][i+1] * (1 - self.config['execution']['slippage'])
                )
                
                if not result:
                    return False
                    
                await asyncio.sleep(0.1)
                
            return True

        except Exception as e:
            logger.error(f"Multi-hop execution error: {e}")
            return False

    async def _execute_direct_trade(self, opportunity: Dict) -> bool:
        """Execute direct swap"""
        try:
            return await self.trader.trading_core.execute_trade(
                pool_address=self.config['accounts']['sol_wsol']['amm_id'],
                amount_in=opportunity['amount_in'],
                min_amount_out=opportunity['amount_out'] * (1 - self.config['execution']['slippage'])
            )

        except Exception as e:
            logger.error(f"Direct trade execution error: {e}")
            return False

    def _calculate_confidence(self, opportunity: Dict, patterns: List[Dict]) -> float:
        """Calculate confidence score based on patterns and metrics"""
        try:
            base_confidence = min(1.0, opportunity['profit'] / self.config['trading']['limits']['max_profit_pct'])
            
            # Adjust based on similar patterns
            if patterns:
                pattern_confidence = max(p['similarity'] for p in patterns)
                pattern_weight = 0.3
                base_confidence = base_confidence * (1 - pattern_weight) + pattern_confidence * pattern_weight

            # Penalize for impact
            impact_penalty = opportunity['impact'] / self.config['trading']['limits']['max_price_impact']
            final_confidence = max(0.0, base_confidence - impact_penalty)

            return final_confidence

        except Exception as e:
            logger.error(f"Confidence calculation error: {e}")
            return 0.0

    async def _optimize_priority_fee(self, opportunity_type: str, profit: float) -> int:
        """Calculate optimal priority fee based on opportunity type and expected profit"""
        try:
            # Get base fee from Helius
            base_fee = await self._get_helius_priority_fee()
            
            # Adjust based on opportunity type and profit
            if opportunity_type == 'sandwich':
                # Sandwich needs higher priority
                fee_multiplier = 1.5
            elif profit > self.config['rpc']['helius']['fee_levels']['profit_threshold']:
                # High profit opportunities justify higher fees
                fee_multiplier = self.config['execution']['fee_multiplier']
            else:
                fee_multiplier = 1.0

            priority_fee = int(base_fee * fee_multiplier)
            
            # Ensure within configured limits
            return min(
                max(priority_fee, self.config['execution']['min_priority_fee']),
                self.config['execution']['max_priority_fee']
            )

        except Exception as e:
            logger.error(f"Priority fee calculation error: {e}")
            return self.config['execution']['priority_fee']

    async def _get_helius_priority_fee(self) -> int:
        """Get current priority fee from Helius API"""
        try:
            async with websockets.connect(self.config['rpc']['helius']['websocket']) as ws:
                await ws.send(json.dumps({
                    "jsonrpc": "2.0",
                    "id": "priority-fee",
                    "method": "getPriorityFeeEstimate",
                    "params": [{
                        "accountKeys": [self.config['programs']['raydium']],
                        "options": {
                            "includeAllPriorityFeeLevels": True
                        }
                    }]
                }))
                
                response = await ws.recv()
                data = json.loads(response)
                
                if 'result' in data:
                    fee_level = self.config['rpc']['helius']['fee_levels']['default']
                    return int(data['result'][fee_level])
                
                return self.config['execution']['priority_fee']

        except Exception as e:
            logger.error(f"Helius priority fee fetch error: {e}")
            return self.config['execution']['priority_fee']

    async def _prepare_transaction(self, opportunity: Dict) -> Dict:
        """Prepare transaction with optimized priority fee"""
        try:
            # Calculate optimal priority fee
            priority_fee = await self._optimize_priority_fee(
                opportunity['type'],
                opportunity['profit']
            )

            # Create transaction parameters
            tx_params = {
                'priority_fee': priority_fee,
                'compute_limit': self.config['execution']['compute_limit'],
                'skip_preflight': True
            }

            if opportunity['type'] == 'sandwich':
                # Add special handling for sandwich transactions
                tx_params.update({
                    'front_run_fee': int(priority_fee * 1.2),  # Higher fee for front-run
                    'back_run_fee': int(priority_fee * 0.8)    # Lower fee for back-run
                })

            return tx_params

        except Exception as e:
            logger.error(f"Transaction preparation error: {e}")
            return {}

    async def _monitor_pool_state(self, pool_address: str):
        """Monitor pool state for optimal execution"""
        try:
            async with websockets.connect(self.config['rpc']['helius']['websocket']) as ws:
                await ws.send(json.dumps({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "accountSubscribe",
                    "params": [
                        pool_address,
                        {
                            "encoding": "jsonParsed",
                            "commitment": "processed"
                        }
                    ]
                }))

                while self.active:
                    msg = await ws.recv()
                    data = json.loads(msg)
                    if 'params' in data:
                        await self._update_pool_state(pool_address, data['params'])

        except Exception as e:
            logger.error(f"Pool monitoring error: {e}")

    async def _update_pool_state(self, pool_address: str, data: Dict):
        """Update pool state with latest data"""
        try:
            with torch.cuda.stream(self.streams['compute']):
                # Update vector store
                await self.trader.vector_store.update_pool_vectors({
                    'address': pool_address,
                    'price': data.get('price', 0),
                    'volume': data.get('volume', 0),
                    'liquidity': data.get('liquidity', 0),
                    'timestamp': time.time()
                })

        except Exception as e:
            logger.error(f"Pool state update error: {e}")

    def _log_metrics(self):
        """Log detailed performance metrics"""
        elapsed_time = time.time() - self.metrics.start_time
        messages_per_second = self.metrics.messages_received / elapsed_time
        opportunities_per_minute = (self.metrics.opportunities_found / elapsed_time) * 60
        trades_per_minute = (self.metrics.successful_trades / elapsed_time) * 60

        logger.info(f"""
Performance Metrics:
Messages/sec: {messages_per_second:.2f}
Opportunities/min: {opportunities_per_minute:.2f}
Trades/min: {trades_per_minute:.2f}
Total Messages: {self.metrics.messages_received}
Total Opportunities: {self.metrics.opportunities_found}
Successful Trades: {self.metrics.successful_trades}
Total Profit: {float(self.metrics.total_profit):.6f} SOL
Sandwich Profits: {float(self.metrics.sandwich_profits):.6f} SOL
Arbitrage Profits: {float(self.metrics.arbitrage_profits):.6f} SOL
        """)

    async def stop(self):
        """Stop advanced trading"""
        logger.info("Stopping advanced trading...")
        self.active = False
        
        # Cleanup resources
        try:
            torch.cuda.empty_cache()
            for stream in self.streams.values():
                stream.synchronize()
            logger.info("GPU resources cleaned up")
        except Exception as e:
            logger.error(f"Cleanup error: {e}")
    async def _safe_execute_trade(self, opportunity: Dict) -> bool:
        """Enhanced trade execution with safety checks"""
        try:
            # Pre-trade validation
            balance = await self.trader.trading_core.get_balance()
            min_balance = self.config['trading']['position']['gas_reserve']
            
            if balance < min_balance:
                logger.warning(f"Insufficient balance: {balance} SOL < {min_balance} SOL required")
                return False

            # Calculate optimal priority fee
            priority_fee = await self._optimize_priority_fee(
                opportunity['type'],
                opportunity['profit']
            )

            # Execute trade with retry logic
            return await self._execute_with_retry(
                lambda: self._execute_opportunity_with_timeout(opportunity, priority_fee),
                max_retries=self.config['execution']['retry_count']
            )

        except Exception as e:
            logger.error(f"Safe trade execution error: {e}")
            return False

    async def _optimize_priority_fee(self, opportunity_type: str, profit: float) -> int:
        """Calculate optimal priority fee based on market conditions"""
        try:
            # Get base fee from Helius
            base_fee = await self._get_helius_priority_fee()
            
            # Adjust based on opportunity type
            multipliers = {
                'sandwich': 1.5,  # Higher priority for sandwich trades
                'triangle': 1.3,  # Medium priority for triangular arb
                'multi_hop': 1.2,  # Slightly higher for multi-hop
                'direct': 1.0     # Base priority for direct swaps
            }
            
            fee_multiplier = multipliers.get(opportunity_type, 1.0)
            
            # Adjust for profit potential
            if profit > self.config['rpc']['helius']['fee_levels']['profit_threshold']:
                fee_multiplier *= self.config['execution']['fee_multiplier']

            priority_fee = int(base_fee * fee_multiplier)
            
            # Ensure within limits
            return min(
                max(priority_fee, self.config['execution']['min_priority_fee']),
                self.config['execution']['max_priority_fee']
            )

        except Exception as e:
            logger.error(f"Priority fee optimization error: {e}")
            return self.config['execution']['priority_fee']

    async def _execute_with_retry(self, execute_func, max_retries: int = 3) -> bool:
        """Execute transaction with exponential backoff retry"""
        retry_count = 0
        base_delay = 0.1  # 100ms initial delay

        while retry_count < max_retries:
            try:
                result = await execute_func()
                if result:
                    return True

                delay = base_delay * (2 ** retry_count)
                logger.warning(f"Transaction failed, retrying in {delay:.2f}s")
                await asyncio.sleep(delay)
                retry_count += 1

            except Exception as e:
                logger.error(f"Execution attempt {retry_count + 1} failed: {e}")
                await asyncio.sleep(base_delay)
                retry_count += 1

        return False

    async def _execute_opportunity_with_timeout(self, opportunity: Dict, priority_fee: int) -> bool:
        """Execute opportunity with timeout protection"""
        try:
            timeout = self.config['system']['websocket'].get('timeout', 5.0)
            
            async with asyncio.timeout(timeout):
                if opportunity['type'] == 'sandwich':
                    return await self._execute_sandwich_with_monitoring(opportunity, priority_fee)
                elif opportunity['type'] == 'triangle':
                    return await self._execute_triangle_with_monitoring(opportunity, priority_fee)
                elif opportunity['type'] == 'multi_hop':
                    return await self._execute_multi_hop_with_monitoring(opportunity, priority_fee)
                else:
                    return await self._execute_direct_with_monitoring(opportunity, priority_fee)

        except asyncio.TimeoutError:
            logger.error(f"Execution timeout for {opportunity['type']} trade")
            return False
        except Exception as e:
            logger.error(f"Execution error: {e}")
            return False

    async def _execute_sandwich_with_monitoring(self, opportunity: Dict, priority_fee: int) -> bool:
        """Execute sandwich trade with monitoring"""
        try:
            # Execute front-run with higher priority
            front_run_fee = int(priority_fee * 1.2)
            front_run_result = await self.trader.trading_core.execute_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['front_run']['amount'],
                min_amount_out=opportunity['front_run']['min_out'],
                priority_fee=front_run_fee
            )

            if not front_run_result:
                logger.error("Front-run failed")
                return False

            # Small delay for block confirmation
            await asyncio.sleep(0.1)

            # Execute back-run with slightly lower priority
            back_run_fee = int(priority_fee * 0.8)
            back_run_result = await self.trader.trading_core.execute_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['back_run']['amount'],
                min_amount_out=opportunity['back_run']['min_out'],
                priority_fee=back_run_fee
            )

            success = bool(back_run_result)
            if success:
                self.metrics.sandwich_profits += Decimal(str(opportunity['profit']))
            
            return success

        except Exception as e:
            logger.error(f"Sandwich execution error: {e}")
            return False

    async def _execute_triangle_with_monitoring(self, opportunity: Dict, priority_fee: int) -> bool:
        """Execute triangle trade with monitoring"""
        try:
            for i, pool in enumerate(opportunity['path']):
                # Adjust fee for each hop
                hop_fee = int(priority_fee * (1.0 - i * 0.1))  # Decrease priority slightly each hop
                
                result = await self.trader.trading_core.execute_trade(
                    pool_address=pool,
                    amount_in=opportunity['amounts'][i],
                    min_amount_out=opportunity['amounts'][i+1] * (1 - self.config['execution']['slippage']),
                    priority_fee=hop_fee
                )
                
                if not result:
                    logger.error(f"Triangle trade failed at hop {i}")
                    return False
                
                await asyncio.sleep(0.1)
            
            self.metrics.arbitrage_profits += Decimal(str(opportunity['profit']))
            return True

        except Exception as e:
            logger.error(f"Triangle execution error: {e}")
            return False

    async def _validate_pool_state(self, pool: str, expected_impact: float) -> bool:
        """Validate pool state before execution"""
        try:
            current_state = await self.trader.trading_core.get_pool_info(pool)
            if not current_state:
                return False

            # Check if price impact would still be acceptable
            impact_buffer = 1.2  # Allow 20% variance
            max_impact = self.config['trading']['limits']['max_price_impact']
            
            return expected_impact * impact_buffer <= max_impact

        except Exception as e:
            logger.error(f"Pool validation error: {e}")
            return False

    async def _update_metrics(self, opportunity: Dict, success: bool):
        """Update trading metrics after execution"""
        try:
            if success:
                self.metrics.successful_trades += 1
                profit = Decimal(str(opportunity['profit']))
                self.metrics.total_profit += profit
                
                # Update strategy-specific metrics
                if opportunity['type'] == 'sandwich':
                    self.metrics.sandwich_profits += profit
                elif opportunity['type'] in ['triangle', 'multi_hop']:
                    self.metrics.arbitrage_profits += profit
                
            else:
                self.metrics.failed_attempts[opportunity['type']] = \
                    self.metrics.failed_attempts.get(opportunity['type'], 0) + 1

        except Exception as e:
            logger.error(f"Metrics update error: {e}")   import asyncio
import base64
import json
import logging
import time
import websockets

from typing import Dict, List, Optional, Set
from decimal import Decimal
from dataclasses import dataclass, field

import torch

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class TradingMetrics:
    messages_received: int = 0
    opportunities_found: int = 0
    successful_trades: int = 0
    total_profit: Decimal = Decimal('0')
    start_time: float = field(default_factory=time.time)
    sandwich_profits: Decimal = Decimal('0')
    arbitrage_profits: Decimal = Decimal('0')
    failed_attempts: Dict[str, int] = field(default_factory=dict)
    active_paths: Set[str] = field(default_factory=set)

class AdvancedTrading:
    def __init__(self, trader, config):
        self.trader = trader
        self.config = config
        self.active = False
        self.metrics = TradingMetrics()

        # Initialize CUDA device and streams if available
        if torch.cuda.is_available():
            self.device = torch.device('cuda')
            self.streams = {
                'arbitrage': torch.cuda.Stream(),
                'sandwich': torch.cuda.Stream(),
                'direct': torch.cuda.Stream(),
                'compute': torch.cuda.Stream()
            }
            self._setup_processing_buffers()
            logger.info("Advanced trading initialized with GPU acceleration")
        else:
            self.device = torch.device('cpu')
            self.streams = {}
            logger.warning("CUDA not available. Running on CPU.")

    def _setup_processing_buffers(self):
        """Initialize GPU processing buffers"""
        try:
            with torch.cuda.stream(self.streams['compute']):
                self.opportunity_buffer = torch.zeros(
                    (1024, 128),
                    device=self.device,
                    dtype=torch.float16
                ).pin_memory()

                self.price_buffer = torch.zeros(
                    (512, 2),
                    device=self.device,
                    dtype=torch.float32
                ).pin_memory()

                logger.debug("GPU buffers initialized")
        except Exception as e:
            logger.error(f"Buffer initialization error: {e}")
            raise

    async def start_websocket_monitoring(self):
        """Start monitoring WebSocket for opportunities"""
        self.active = True
        logger.info("Starting advanced trading monitor...")

        while self.active:
            try:
                async with websockets.connect(
                    self.config['rpc']['helius']['websocket'],
                    ping_interval=self.config['system']['websocket']['ping_interval'],
                    ping_timeout=self.config['system']['websocket']['ping_timeout'],
                    close_timeout=self.config['system']['websocket']['close_timeout'],
                    max_size=self.config['system']['websocket']['max_message_size']
                ) as ws:
                    await self._subscribe_to_logs(ws)

                    while self.active:
                        msg = await ws.recv()
                        self.metrics.messages_received += 1

                        # Process opportunities
                        await self._process_message(msg)

                        # Log metrics periodically
                        if self.metrics.messages_received % 100 == 0:
                            self._log_metrics()

            except websockets.exceptions.ConnectionClosedError:
                logger.warning("WebSocket connection closed. Reconnecting in 1 second...")
                await asyncio.sleep(1)
            except Exception as e:
                logger.error(f"WebSocket error: {e}. Reconnecting in 1 second...")
                await asyncio.sleep(1)

    async def _subscribe_to_logs(self, ws):
        """Subscribe to logs via WebSocket"""
        try:
            subscribe_message = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "logsSubscribe",
                "params": [
                    {"mentions": [self.config['programs']['raydium']]},
                    {"commitment": "processed"}
                ]
            }
            await ws.send(json.dumps(subscribe_message))
            logger.debug("Subscribed to logs")
        except Exception as e:
            logger.error(f"Subscription error: {e}")

    async def _process_message(self, msg: str):
        """Process incoming WebSocket message with strategy checking"""
        try:
            data = json.loads(msg)
            if 'params' not in data:
                return

            logs = data['params'].get('result', {}).get('value', {}).get('logs', [])
            for log in logs:
                if "Program log:" not in log:
                    continue

                # Extract and decode the base64 data
                encoded_data = log.split("Program log: ")[1].strip()
                raw_data = base64.b64decode(encoded_data)
                self.metrics.opportunities_found += 1

                # Check all strategies in parallel
                results = await asyncio.gather(
                    self._check_sandwich_opportunity(raw_data),
                    self._check_triangle_arbitrage(raw_data),
                    self._check_multi_hop_opportunity(raw_data),
                    self._check_direct_swap(raw_data),
                    return_exceptions=True
                )

                valid_opportunities = [
                    r for r in results if isinstance(r, dict) and r.get('profit', 0) > 0
                ]

                if valid_opportunities:
                    # Sort by profit/risk ratio
                    best_opp = max(
                        valid_opportunities,
                        key=lambda x: x['profit'] * x.get('confidence', 1.0) / x.get('impact', 1.0)
                    )

                    logger.info(f"Found opportunity: {best_opp['type']} with profit {best_opp['profit']:.4f}%")
                    await self._execute_opportunity(best_opp)

        except json.JSONDecodeError:
            logger.error("JSON decode error in message processing")
        except Exception as e:
            logger.error(f"Message processing error: {e}")

    async def _check_sandwich_opportunity(self, raw_data: bytes) -> Optional[Dict]:
        """Check for sandwich opportunities"""
        try:
            # Decode transaction data
            victim_tx = await self._decode_victim_transaction(raw_data)
            if not victim_tx:
                return None

            pool_address = victim_tx['pool']

            # Simulate front-run and back-run trades
            front_run = await self._simulate_front_run(victim_tx)
            back_run = await self._simulate_back_run(victim_tx, front_run)

            if not front_run or not back_run:
                return None

            # Calculate expected profit
            profit = back_run['amount_out'] - front_run['amount_in']

            if profit <= 0:
                return None

            return {
                'type': 'sandwich',
                'profit': float(profit),
                'front_run': front_run,
                'back_run': back_run,
                'victim': victim_tx,
                'pool': pool_address,
                'confidence': 0.9,  # High confidence for sandwich
                'impact': front_run['price_impact']
            }

        except Exception as e:
            logger.error(f"Sandwich check error: {e}")
            return None

    async def _decode_victim_transaction(self, raw_data: bytes) -> Optional[Dict]:
        """Decode the victim transaction from raw data"""
        try:
            # For demonstration, assume raw_data is JSON-encoded
            tx_data = json.loads(raw_data.decode('utf-8'))
            return tx_data
        except Exception as e:
            logger.error(f"Error decoding victim transaction: {e}")
            return None

    async def _simulate_front_run(self, victim_tx: Dict) -> Optional[Dict]:
        """Simulate the front-run trade"""
        try:
            # Placeholder for simulation logic
            amount_in = victim_tx['amount']
            # Simulate price impact and amount_out
            price_impact = 0.005  # Assume 0.5% price impact
            amount_out = amount_in * (1 - price_impact)

            return {
                'amount_in': amount_in,
                'amount_out': amount_out,
                'price_impact': price_impact,
                'min_out': amount_out * (1 - self.config['execution']['slippage'])
            }
        except Exception as e:
            logger.error(f"Front-run simulation error: {e}")
            return None

    async def _simulate_back_run(self, victim_tx: Dict, front_run: Dict) -> Optional[Dict]:
        """Simulate the back-run trade"""
        try:
            # Placeholder for simulation logic
            amount_in = front_run['amount_out']
            price_impact = 0.005  # Assume 0.5% price impact
            amount_out = amount_in * (1 + price_impact)

            return {
                'amount_in': amount_in,
                'amount_out': amount_out,
                'price_impact': price_impact,
                'min_out': amount_out * (1 - self.config['execution']['slippage'])
            }
        except Exception as e:
            logger.error(f"Back-run simulation error: {e}")
            return None

    async def _check_triangle_arbitrage(self, raw_data: bytes) -> Optional[Dict]:
        """Check for triangle arbitrage opportunities"""
        try:
            # Decode token information
            token_info = await self._decode_token_info(raw_data)
            if not token_info:
                return None

            # Find arbitrage paths
            paths = await self.trader.arbitrage_finder.find_triangle_paths(token_info['token_in'])
            if not paths:
                return None

            best_profit = 0
            best_path = None
            best_amounts = None
            best_impact = None

            for path in paths:
                amounts, impact, profit = await self._simulate_arbitrage_path(path)
                if profit > best_profit:
                    best_profit = profit
                    best_path = path
                    best_amounts = amounts
                    best_impact = impact

            if best_profit <= 0:
                return None

            return {
                'type': 'triangle',
                'profit': float(best_profit),
                'path': best_path,
                'amounts': best_amounts,
                'impact': best_impact,
                'confidence': 0.8
            }

        except Exception as e:
            logger.error(f"Triangle arbitrage error: {e}")
            return None

    async def _simulate_arbitrage_path(self, path: List[str]) -> (List[float], float, float):
        """Simulate the arbitrage path to estimate profit"""
        try:
            amounts = [self.config['trading']['position']['size']]
            impact = 0
            for i in range(len(path) - 1):
                # Simulate swap
                amount_in = amounts[-1]
                amount_out = amount_in * 0.99  # Assume 1% fee
                amounts.append(amount_out)
                impact += 0.01  # Accumulate price impact

            profit = amounts[-1] - amounts[0]
            return amounts, impact, profit
        except Exception as e:
            logger.error(f"Arbitrage path simulation error: {e}")
            return [], 0, 0

    async def _decode_token_info(self, raw_data: bytes) -> Optional[Dict]:
        """Decode token information from raw data"""
        try:
            # For demonstration, assume raw_data is JSON-encoded
            token_data = json.loads(raw_data.decode('utf-8'))
            return token_data
        except Exception as e:
            logger.error(f"Error decoding token info: {e}")
            return None

    async def _check_multi_hop_opportunity(self, raw_data: bytes) -> Optional[Dict]:
        """Check for multi-hop opportunities"""
        try:
            # Decode token information
            token_info = await self._decode_token_info(raw_data)
            if not token_info:
                return None

            # Find optimal route
            optimal_route = await self.trader.route_optimizer.find_best_route(
                self.config['trading']['position']['size'],
                token_info['token_in'],
                token_info['token_out']
            )

            if not optimal_route or optimal_route['profit'] <= 0:
                return None

            return {
                'type': 'multi_hop',
                'profit': float(optimal_route['profit']),
                'path': optimal_route['path'],
                'amounts': optimal_route['amounts'],
                'impact': optimal_route['impact'],
                'confidence': 0.75
            }

        except Exception as e:
            logger.error(f"Multi-hop check error: {e}")
            return None

    async def _check_direct_swap(self, raw_data: bytes) -> Optional[Dict]:
        """Check for direct swap opportunities"""
        try:
            # Decode market data
            market_data = await self._decode_market_data(raw_data)
            if not market_data:
                return None

            # Simulate direct swap
            amount_in = self.config['trading']['position']['size']
            amount_out = amount_in * market_data['price']
            impact = market_data['price_impact']

            profit = amount_out - amount_in
            if profit <= 0:
                return None

            return {
                'type': 'direct',
                'profit': float(profit),
                'amount_in': amount_in,
                'amount_out': amount_out,
                'impact': impact,
                'confidence': 0.7
            }

        except Exception as e:
            logger.error(f"Direct swap check error: {e}")
            return None

    async def _decode_market_data(self, raw_data: bytes) -> Optional[Dict]:
        """Decode market data from raw data"""
        try:
            # For demonstration, assume raw_data is JSON-encoded
            market_data = json.loads(raw_data.decode('utf-8'))
            return market_data
        except Exception as e:
            logger.error(f"Error decoding market data: {e}")
            return None

    async def _execute_opportunity(self, opportunity: Dict) -> bool:
        """Execute trading opportunity"""
        try:
            logger.info(f"Executing {opportunity['type']} opportunity...")

            if opportunity['type'] == 'sandwich':
                result = await self._execute_sandwich_trade(opportunity)
            elif opportunity['type'] == 'triangle':
                result = await self._execute_triangle_trade(opportunity)
            elif opportunity['type'] == 'multi_hop':
                result = await self._execute_multi_hop_trade(opportunity)
            else:  # direct swap
                result = await self._execute_direct_trade(opportunity)

            await self._update_metrics(opportunity, result)
            return result

        except Exception as e:
            logger.error(f"Trade execution error: {e}")
            return False

    async def _execute_sandwich_trade(self, opportunity: Dict) -> bool:
        """Execute sandwich trade"""
        try:
            # Prepare transactions
            front_run_tx = await self.trader.trading_core.prepare_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['front_run']['amount_in'],
                min_amount_out=opportunity['front_run']['min_out']
            )

            back_run_tx = await self.trader.trading_core.prepare_trade(
                pool_address=opportunity['pool'],
                amount_in=opportunity['back_run']['amount_in'],
                min_amount_out=opportunity['back_run']['min_out']
            )

            # Send transactions with appropriate timing
            await self.trader.trading_core.send_transaction(front_run_tx)
            # Wait for victim transaction to be processed (implementation-specific)
            await asyncio.sleep(0.5)
            await self.trader
 can you create the best script from the code i provided please