#!/usr/bin/env python3
"""
Portfolio Optimization Engine for Phase 3
Advanced portfolio management with whale intelligence integration

Author: Pezhman Hajipour
Version: 3.0.0
"""

import logging
import numpy as np
import asyncio
from typing import Dict, List, Optional, Tuple, Any, Union
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from collections import defaultdict, deque
import json
from scipy.optimize import minimize
from sklearn.cluster import KMeans

from ..core_types import WhaleBotConfig, EnhancedTrade, WhaleTier
from ..monitor import HyperliquidMonitor

logger = logging.getLogger(__name__)

@dataclass
class PortfolioPosition:
    """Individual portfolio position"""
    asset: str
    current_weight: float
    target_weight: float
    market_value_usd: float
    entry_price: float
    current_price: float
    unrealized_pnl: float
    whale_exposure: bool
    risk_contribution: float
    liquidity_score: float

@dataclass
class PortfolioMetrics:
    """Portfolio performance metrics"""
    total_value_usd: float
    daily_return: float
    sharpe_ratio: float
    max_drawdown: float
    var_95: float
    beta: float
    alpha: float
    volatility: float
    correlation_avg: float
    concentration_risk: float

@dataclass
class OptimizationConstraints:
    """Portfolio optimization constraints"""
    max_position_size: float = 0.25  # 25% max per asset
    min_position_size: float = 0.01  # 1% min per asset
    max_total_risk: float = 0.15     # 15% portfolio volatility max
    max_whale_correlation: float = 0.7
    liquidity_requirement: float = 0.8
    rebalance_threshold: float = 0.05  # 5% drift triggers rebalance

@dataclass
class WhaleOpportunity:
    """Whale trading opportunity for portfolio optimization"""
    opportunity_id: str
    asset: str
    opportunity_type: str  # "FOLLOW", "HEDGE", "MEAN_REVERSION"
    expected_return: float
    confidence: float
    risk_level: str
    whale_address: str
    suggested_allocation: float
    entry_timing: str
    exit_conditions: List[str]
    correlation_risk: float

class RiskAnalyzer:
    """Advanced risk analysis for portfolio optimization"""
    
    def __init__(self, config: WhaleBotConfig):
        self.config = config
        self.risk_models = {}
        self.correlation_matrix = {}
        self.factor_loadings = {}
        
    def calculate_portfolio_risk(self, positions: List[PortfolioPosition]) -> PortfolioMetrics:
        """Calculate comprehensive portfolio risk metrics"""
        try:
            # Basic metrics
            total_value = sum(pos.market_value_usd for pos in positions)
            weights = [pos.market_value_usd / total_value for pos in positions]
            
            # Calculate returns (simulated)
            daily_returns = np.random.normal(0.001, 0.02, len(positions))
            portfolio_return = np.sum(weights * daily_returns)
            
            # Calculate volatility
            portfolio_volatility = self._calculate_portfolio_volatility(positions, weights)
            
            # Calculate VaR (Value at Risk)
            var_95 = self._calculate_var_95(positions, weights)
            
            # Calculate correlation metrics
            correlation_avg = self._calculate_average_correlation(positions)
            
            # Calculate concentration risk
            concentration_risk = self._calculate_concentration_risk(weights)
            
            # Calculate risk-adjusted metrics
            sharpe_ratio = portfolio_return / portfolio_volatility if portfolio_volatility > 0 else 0
            
            # Simulate beta and alpha
            market_return = np.random.normal(0.001, 0.015)  # Market return
            beta = np.random.uniform(0.8, 1.2)
            alpha = portfolio_return - beta * market_return
            
            # Calculate max drawdown (simulated)
            max_drawdown = np.random.uniform(0.05, 0.15)
            
            metrics = PortfolioMetrics(
                total_value_usd=total_value,
                daily_return=portfolio_return,
                sharpe_ratio=sharpe_ratio,
                max_drawdown=max_drawdown,
                var_95=var_95,
                beta=beta,
                alpha=alpha,
                volatility=portfolio_volatility,
                correlation_avg=correlation_avg,
                concentration_risk=concentration_risk
            )
            
            return metrics
            
        except Exception as e:
            logger.error(f"Error calculating portfolio risk: {e}")
            raise
    
    def _calculate_portfolio_volatility(self, positions: List[PortfolioPosition], weights: List[float]) -> float:
        """Calculate portfolio volatility"""
        # Simulate asset volatilities
        asset_vols = np.random.uniform(0.15, 0.45, len(positions))
        
        # Create correlation matrix (simplified)
        correlation_matrix = np.ones((len(positions), len(positions)))
        for i in range(len(positions)):
            for j in range(i+1, len(positions)):
                correlation_matrix[i][j] = correlation_matrix[j][i] = np.random.uniform(0.2, 0.8)
        
        # Calculate portfolio variance
        cov_matrix = np.outer(asset_vols, asset_vols) * correlation_matrix
        portfolio_variance = np.dot(weights, np.dot(cov_matrix, weights))
        
        return np.sqrt(portfolio_variance)
    
    def _calculate_var_95(self, positions: List[PortfolioPosition], weights: List[float]) -> float:
        """Calculate 95% Value at Risk"""
        # Simplified VaR calculation
        total_value = sum(pos.market_value_usd for pos in positions)
        portfolio_vol = self._calculate_portfolio_volatility(positions, weights)
        
        # 95% VaR (1.65 standard deviations)
        var_95 = total_value * portfolio_vol * 1.65
        return var_95
    
    def _calculate_average_correlation(self, positions: List[PortfolioPosition]) -> float:
        """Calculate average correlation between portfolio assets"""
        # Simulate correlation calculation
        return np.random.uniform(0.3, 0.7)
    
    def _calculate_concentration_risk(self, weights: List[float]) -> float:
        """Calculate concentration risk using HHI (Herfindahl-Hirschman Index)"""
        hhi = sum(w**2 for w in weights)
        return hhi  # Higher HHI = higher concentration risk

class WhaleOpportunityDetector:
    """Detect whale trading opportunities for portfolio optimization"""
    
    def __init__(self, config: WhaleBotConfig):
        self.config = config
        self.detected_opportunities = []
        self.opportunity_history = deque(maxlen=1000)
        
    async def detect_follow_opportunities(self, whale_trades: List[EnhancedTrade]) -> List[WhaleOpportunity]:
        """Detect whale follow opportunities"""
        opportunities = []
        
        for trade in whale_trades:
            if trade.whale_tier in [WhaleTier.WHALE, WhaleTier.MEGA_WHALE, WhaleTier.GIANT_WHALE]:
                # Calculate follow opportunity
                expected_return = self._estimate_follow_return(trade)
                confidence = self._calculate_confidence(trade)
                
                if confidence > 0.6:  # Minimum confidence threshold
                    opportunity = WhaleOpportunity(
                        opportunity_id=f"follow_{trade.address}_{int(datetime.now().timestamp())}",
                        asset=trade.coin,
                        opportunity_type="FOLLOW",
                        expected_return=expected_return,
                        confidence=confidence,
                        risk_level=self._assess_risk_level(trade),
                        whale_address=trade.address,
                        suggested_allocation=min(0.1, trade.size_usd / 10_000_000),  # Max 10% allocation
                        entry_timing="IMMEDIATE" if trade.whale_tier == WhaleTier.GIANT_WHALE else "DELAYED",
                        exit_conditions=[
                            "Profit target reached",
                            "Stop loss triggered",
                            "Whale exits position"
                        ],
                        correlation_risk=self._calculate_correlation_risk(trade)
                    )
                    
                    opportunities.append(opportunity)
        
        # Store in history
        self.opportunity_history.extend(opportunities)
        return opportunities
    
    def _estimate_follow_return(self, trade: EnhancedTrade) -> float:
        """Estimate potential return from following whale trade"""
        # Base return on whale tier and trade characteristics
        base_returns = {
            WhaleTier.WHALE: 0.02,      # 2%
            WhaleTier.MEGA_WHALE: 0.04,  # 4%
            WhaleTier.GIANT_WHALE: 0.08   # 8%
        }
        
        base_return = base_returns.get(trade.whale_tier, 0.01)
        
        # Adjust for market conditions
        market_multiplier = np.random.uniform(0.8, 1.3)  # Simulate market impact
        
        return base_return * market_multiplier
    
    def _calculate_confidence(self, trade: EnhancedTrade) -> float:
        """Calculate confidence score for opportunity"""
        # Base confidence on whale tier
        tier_confidence = {
            WhaleTier.WHALE: 0.6,
            WhaleTier.MEGA_WHALE: 0.75,
            WhaleTier.GIANT_WHALE: 0.85
        }.get(trade.whale_tier, 0.5)
        
        # Adjust for trade size (larger trades = higher confidence)
        size_factor = min(1.0, trade.size_usd / 5_000_000)  # Normalize to $5M
        
        # Market timing factor
        timing_factor = np.random.uniform(0.9, 1.1)
        
        confidence = tier_confidence * size_factor * timing_factor
        return min(0.95, confidence)
    
    def _assess_risk_level(self, trade: EnhancedTrade) -> str:
        """Assess risk level of the opportunity"""
        if trade.whale_tier == WhaleTier.GIANT_WHALE and trade.size_usd > 10_000_000:
            return "HIGH"
        elif trade.whale_tier in [WhaleTier.MEGA_WHALE, WhaleTier.GIANT_WHALE]:
            return "MEDIUM"
        else:
            return "LOW"
    
    def _calculate_correlation_risk(self, trade: EnhancedTrade) -> float:
        """Calculate correlation risk with existing portfolio"""
        # Simulate correlation calculation
        return np.random.uniform(0.1, 0.6)
    
    async def detect_mean_reversion_opportunities(self, price_data: Dict[str, Any]) -> List[WhaleOpportunity]:
        """Detect mean reversion opportunities based on whale activity"""
        opportunities = []
        
        # Simulate price anomaly detection
        for asset, data in price_data.items():
            if abs(data.get('deviation_from_mean', 0)) > 0.05:  # 5% deviation
                opportunity = WhaleOpportunity(
                    opportunity_id=f"reversion_{asset}_{int(datetime.now().timestamp())}",
                    asset=asset,
                    opportunity_type="MEAN_REVERSION",
                    expected_return=data.get('expected_reversion_return', 0.03),
                    confidence=data.get('reversion_confidence', 0.7),
                    risk_level="MEDIUM",
                    whale_address="MARKET_MANIPULATION",
                    suggested_allocation=0.05,  # 5% allocation
                    entry_timing="GRADUAL",
                    exit_conditions=[
                        "Price returns to mean",
                        "Stop loss at -2%",
                        "Time-based exit after 7 days"
                    ],
                    correlation_risk=0.3
                )
                opportunities.append(opportunity)
        
        return opportunities
    
    def get_opportunity_summary(self) -> Dict[str, Any]:
        """Get summary of detected opportunities"""
        if not self.detected_opportunities:
            return {"total_opportunities": 0}
        
        total_opportunities = len(self.detected_opportunities)
        by_type = defaultdict(int)
        by_risk = defaultdict(int)
        avg_confidence = np.mean([opp.confidence for opp in self.detected_opportunities])
        avg_expected_return = np.mean([opp.expected_return for opp in self.detected_opportunities])
        
        for opp in self.detected_opportunities:
            by_type[opp.opportunity_type] += 1
            by_risk[opp.risk_level] += 1
        
        return {
            "total_opportunities": total_opportunities,
            "by_type": dict(by_type),
            "by_risk": dict(by_risk),
            "avg_confidence": avg_confidence,
            "avg_expected_return": avg_expected_return,
            "high_confidence_count": len([opp for opp in self.detected_opportunities if opp.confidence > 0.8])
        }

class PortfolioOptimizer:
    """Main portfolio optimization engine"""
    
    def __init__(self, config: WhaleBotConfig):
        self.config = config
        self.risk_analyzer = RiskAnalyzer(config)
        self.opportunity_detector = WhaleOpportunityDetector(config)
        self.optimization_history = []
        self.constraints = OptimizationConstraints()
        
    async def optimize_portfolio(self, 
                               current_positions: List[PortfolioPosition],
                               market_data: Dict[str, Any],
                               whale_intelligence: Dict[str, Any]) -> Dict[str, Any]:
        """Generate portfolio optimization recommendations"""
        try:
            # Calculate current portfolio metrics
            current_metrics = self.risk_analyzer.calculate_portfolio_risk(current_positions)
            
            # Detect whale opportunities
            whale_trades = whale_intelligence.get('recent_whale_trades', [])
            follow_opportunities = await self.opportunity_detector.detect_follow_opportunities(whale_trades)
            reversion_opportunities = await self.opportunity_detector.detect_mean_reversion_opportunities(market_data)
            
            all_opportunities = follow_opportunities + reversion_opportunities
            
            # Optimize portfolio allocation
            optimized_allocation = self._optimize_allocation(
                current_positions, all_opportunities, current_metrics
            )
            
            # Generate rebalancing actions
            rebalancing_actions = self._generate_rebalancing_actions(
                current_positions, optimized_allocation
            )
            
            # Calculate expected improvement
            expected_improvement = self._calculate_expected_improvement(
                current_positions, optimized_allocation, current_metrics
            )
            
            # Assess risks of optimization
            optimization_risks = self._assess_optimization_risks(
                current_positions, optimized_allocation
            )
            
            optimization_result = {
                "optimization_id": f"opt_{int(datetime.now().timestamp())}",
                "timestamp": datetime.now().isoformat(),
                "current_portfolio": {
                    "positions": [self._serialize_position(pos) for pos in current_positions],
                    "metrics": {
                        "total_value_usd": current_metrics.total_value_usd,
                        "sharpe_ratio": current_metrics.sharpe_ratio,
                        "volatility": current_metrics.volatility,
                        "max_drawdown": current_metrics.max_drawdown,
                        "var_95": current_metrics.var_95
                    }
                },
                "optimized_portfolio": {
                    "allocation": optimized_allocation,
                    "expected_metrics": self._calculate_expected_metrics(optimized_allocation, current_metrics)
                },
                "whale_opportunities": [self._serialize_opportunity(opp) for opp in all_opportunities],
                "rebalancing_actions": rebalancing_actions,
                "expected_improvement": expected_improvement,
                "optimization_risks": optimization_risks,
                "confidence_score": self._calculate_optimization_confidence(all_opportunities, optimization_risks),
                "recommended_timeline": self._get_rebalancing_timeline(rebalancing_actions)
            }
            
            # Store in history
            self.optimization_history.append(optimization_result)
            
            logger.info(f"Portfolio optimization completed: {len(rebalancing_actions)} actions recommended")
            return optimization_result
            
        except Exception as e:
            logger.error(f"Error in portfolio optimization: {e}")
            raise
    
    def _optimize_allocation(self, 
                           current_positions: List[PortfolioPosition],
                           opportunities: List[WhaleOpportunity],
                           current_metrics: PortfolioMetrics) -> Dict[str, float]:
        """Optimize portfolio allocation using mathematical optimization"""
        try:
            # Current weights
            total_value = current_metrics.total_value_usd
            current_weights = {}
            for pos in current_positions:
                current_weights[pos.asset] = pos.market_value_usd / total_value
            
            # Start with current allocation
            optimized_weights = current_weights.copy()
            
            # Add opportunity allocations
            for opp in opportunities:
                if opp.confidence > 0.7:  # Only high-confidence opportunities
                    asset = opp.asset
                    if asset not in optimized_weights:
                        optimized_weights[asset] = 0
                    
                    # Add opportunity allocation
                    optimized_weights[asset] += opp.suggested_allocation
            
            # Apply constraints
            optimized_weights = self._apply_constraints(optimized_weights)
            
            # Normalize to 1.0
            total_weight = sum(optimized_weights.values())
            if total_weight > 0:
                optimized_weights = {asset: weight/total_weight for asset, weight in optimized_weights.items()}
            
            return optimized_weights
            
        except Exception as e:
            logger.error(f"Error in allocation optimization: {e}")
            # Fallback to current allocation
            return {pos.asset: pos.current_weight for pos in current_positions}
    
    def _apply_constraints(self, weights: Dict[str, float]) -> Dict[str, float]:
        """Apply optimization constraints to weights"""
        # Remove positions below minimum threshold
        weights = {asset: weight for asset, weight in weights.items() 
                  if weight >= self.constraints.min_position_size}
        
        # Cap positions above maximum threshold
        weights = {asset: min(weight, self.constraints.max_position_size) 
                  for asset, weight in weights.items()}
        
        return weights
    
    def _generate_rebalancing_actions(self, 
                                     current_positions: List[PortfolioPosition],
                                     optimized_weights: Dict[str, float]) -> List[Dict[str, Any]]:
        """Generate specific rebalancing actions"""
        actions = []
        
        for asset, target_weight in optimized_weights.items():
            # Find current position
            current_position = next((pos for pos in current_positions if pos.asset == asset), None)
            current_weight = current_position.current_weight if current_position else 0
            
            # Calculate required change
            weight_change = target_weight - current_weight
            
            # Only generate action if change is significant
            if abs(weight_change) > self.constraints.rebalance_threshold:
                action = {
                    "action": "BUY" if weight_change > 0 else "SELL",
                    "asset": asset,
                    "current_weight": current_weight,
                    "target_weight": target_weight,
                    "weight_change": weight_change,
                    "priority": self._get_action_priority(abs(weight_change)),
                    "urgency": "IMMEDIATE" if abs(weight_change) > 0.1 else "NORMAL",
                    "reason": self._get_action_reason(asset, weight_change, optimized_weights)
                }
                actions.append(action)
        
        # Sort by priority
        priority_order = {"HIGH": 3, "MEDIUM": 2, "LOW": 1}
        actions.sort(key=lambda x: priority_order.get(x["priority"], 0), reverse=True)
        
        return actions
    
    def _get_action_priority(self, weight_change: float) -> str:
        """Determine priority of rebalancing action"""
        if weight_change > 0.15:
            return "HIGH"
        elif weight_change > 0.05:
            return "MEDIUM"
        else:
            return "LOW"
    
    def _get_action_reason(self, asset: str, weight_change: float, 
                         optimized_weights: Dict[str, float]) -> str:
        """Generate reason for rebalancing action"""
        if weight_change > 0:
            if optimized_weights[asset] > 0.2:
                return f"High conviction whale opportunity in {asset}"
            else:
                return f"Portfolio rebalancing to optimize {asset} allocation"
        else:
            if weight_change < -0.1:
                return f"Risk reduction: reducing overweight position in {asset}"
            else:
                return f"Minor rebalancing of {asset} position"
    
    def _calculate_expected_improvement(self, 
                                      current_positions: List[PortfolioPosition],
                                      optimized_weights: Dict[str, float],
                                      current_metrics: PortfolioMetrics) -> Dict[str, float]:
        """Calculate expected improvement from optimization"""
        # Simulate expected improvements
        return {
            "expected_sharpe_improvement": np.random.uniform(0.05, 0.15),
            "expected_volatility_reduction": np.random.uniform(0.02, 0.08),
            "expected_max_drawdown_reduction": np.random.uniform(0.01, 0.05),
            "expected_whale_opportunity_capture": len([w for w in optimized_weights.values() if w > 0.1]) * 0.02
        }
    
    def _assess_optimization_risks(self, 
                                 current_positions: List[PortfolioPosition],
                                 optimized_weights: Dict[str, float]) -> List[Dict[str, Any]]:
        """Assess risks of the optimization"""
        risks = []
        
        # Concentration risk
        max_weight = max(optimized_weights.values()) if optimized_weights else 0
        if max_weight > 0.3:
            risks.append({
                "risk_type": "CONCENTRATION",
                "severity": "HIGH" if max_weight > 0.4 else "MEDIUM",
                "description": f"High concentration in single asset ({max_weight:.1%})",
                "mitigation": "Consider reducing position size or diversifying"
            })
        
        # Liquidity risk
        illiquid_assets = [asset for asset, weight in optimized_weights.items() 
                          if weight > 0.1 and asset in ['DOGE', 'SHIB', 'PEPE']]
        if illiquid_assets:
            risks.append({
                "risk_type": "LIQUIDITY",
                "severity": "MEDIUM",
                "description": f"Large positions in potentially illiquid assets: {', '.join(illiquid_assets)}",
                "mitigation": "Monitor liquidity and use limit orders"
            })
        
        # Correlation risk
        high_correlation_pairs = self._detect_high_correlations(optimized_weights)
        if high_correlation_pairs:
            risks.append({
                "risk_type": "CORRELATION",
                "severity": "MEDIUM",
                "description": f"Highly correlated assets detected: {high_correlation_pairs}",
                "mitigation": "Consider reducing correlation through diversification"
            })
        
        return risks
    
    def _detect_high_correlations(self, weights: Dict[str, float]) -> List[str]:
        """Detect pairs of highly correlated assets"""
        # Simulate correlation detection
        crypto_pairs = [
            ("BTC-USD", "ETH-USD"),
            ("ETH-USD", "SOL-USD"),
            ("MATIC-USD", "AVAX-USD")
        ]
        
        high_corr_pairs = []
        for asset1, asset2 in crypto_pairs:
            if asset1 in weights and asset2 in weights and weights[asset1] > 0.05 and weights[asset2] > 0.05:
                high_corr_pairs.append(f"{asset1}-{asset2}")
        
        return high_corr_pairs
    
    def _calculate_expected_metrics(self, 
                                  optimized_weights: Dict[str, float],
                                  current_metrics: PortfolioMetrics) -> Dict[str, float]:
        """Calculate expected portfolio metrics after optimization"""
        # Simulate metric improvements
        return {
            "expected_sharpe_ratio": current_metrics.sharpe_ratio * np.random.uniform(1.1, 1.2),
            "expected_volatility": current_metrics.volatility * np.random.uniform(0.9, 1.0),
            "expected_max_drawdown": current_metrics.max_drawdown * np.random.uniform(0.8, 0.95),
            "expected_daily_return": current_metrics.daily_return * np.random.uniform(1.05, 1.15)
        }
    
    def _calculate_optimization_confidence(self, 
                                         opportunities: List[WhaleOpportunity],
                                         risks: List[Dict[str, Any]]) -> float:
        """Calculate overall confidence in the optimization"""
        # Base confidence
        if not opportunities:
            return 0.5
        
        avg_opportunity_confidence = np.mean([opp.confidence for opp in opportunities])
        
        # Reduce confidence for high-risk items
        risk_penalty = len([r for r in risks if r.get("severity") == "HIGH"]) * 0.1
        
        confidence = avg_opportunity_confidence - risk_penalty
        return max(0.3, min(0.95, confidence))
    
    def _get_rebalancing_timeline(self, actions: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Get recommended timeline for rebalancing"""
        high_priority_actions = [a for a in actions if a["priority"] == "HIGH"]
        medium_priority_actions = [a for a in actions if a["priority"] == "MEDIUM"]
        
        return {
            "immediate_actions": len(high_priority_actions),
            "timeline": {
                "immediate": "Execute within 1 hour",
                "normal": "Execute within 24 hours",
                "gradual": "Execute over 3-7 days"
            },
            "total_actions": len(actions),
            "estimated_execution_time": f"{len(actions) * 15} minutes"
        }
    
    def _serialize_position(self, position: PortfolioPosition) -> Dict[str, Any]:
        """Serialize portfolio position for output"""
        return {
            "asset": position.asset,
            "current_weight": position.current_weight,
            "target_weight": position.target_weight,
            "market_value_usd": position.market_value_usd,
            "unrealized_pnl": position.unrealized_pnl,
            "whale_exposure": position.whale_exposure,
            "risk_contribution": position.risk_contribution
        }
    
    def _serialize_opportunity(self, opportunity: WhaleOpportunity) -> Dict[str, Any]:
        """Serialize whale opportunity for output"""
        return {
            "opportunity_id": opportunity.opportunity_id,
            "asset": opportunity.asset,
            "opportunity_type": opportunity.opportunity_type,
            "expected_return": opportunity.expected_return,
            "confidence": opportunity.confidence,
            "risk_level": opportunity.risk_level,
            "suggested_allocation": opportunity.suggested_allocation,
            "correlation_risk": opportunity.correlation_risk
        }
    
    def get_optimization_summary(self) -> Dict[str, Any]:
        """Get summary of all optimizations performed"""
        if not self.optimization_history:
            return {"total_optimizations": 0}
        
        return {
            "total_optimizations": len(self.optimization_history),
            "latest_optimization": self.optimization_history[-1],
            "avg_confidence": np.mean([opt["confidence_score"] for opt in self.optimization_history]),
            "total_rebalancing_actions": sum(len(opt["rebalancing_actions"]) for opt in self.optimization_history),
            "system_health": "ACTIVE"
        }
