import numpy as np
import pandas as pd
import cvxpy as cp
from typing import Dict, List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)

class PortfolioOptimizer:
    """
    Portfolio optimization using various methods
    """
    
    def __init__(self, risk_free_rate: float = 0.02):
        self.risk_free_rate = risk_free_rate
        
    def optimize(
        self,
        returns: np.ndarray,
        method: str = 'mean_variance',
        constraints: Optional[Dict] = None
    ) -> Dict:
        """
        Optimize portfolio weights using specified method
        """
        if method == 'mean_variance':
            return self._mean_variance_optimization(returns, constraints)
        elif method == 'min_variance':
            return self._minimum_variance_optimization(returns, constraints)
        elif method == 'max_sharpe':
            return self._maximum_sharpe_optimization(returns, constraints)
        elif method == 'risk_parity':
            return self._risk_parity_optimization(returns)
        elif method == 'black_litterman':
            return self._black_litterman_optimization(returns, constraints)
        else:
            raise ValueError(f"Unknown optimization method: {method}")
    
    def _mean_variance_optimization(
        self,
        returns: np.ndarray,
        constraints: Optional[Dict] = None
    ) -> Dict:
        """
        Classical Markowitz mean-variance optimization
        """
        n_assets = returns.shape[1]
        
        # Calculate expected returns and covariance
        expected_returns = np.mean(returns, axis=0)
        cov_matrix = np.cov(returns.T)
        
        # Add small value to diagonal for numerical stability
        cov_matrix += np.eye(n_assets) * 1e-6
        
        # Define optimization variables
        weights = cp.Variable(n_assets)
        
        # Define objective (maximize return - lambda * risk)
        risk_aversion = 2.0  # Can be parameterized
        portfolio_return = expected_returns @ weights
        portfolio_risk = cp.quad_form(weights, cov_matrix)
        
        objective = cp.Maximize(portfolio_return - risk_aversion * portfolio_risk)
        
        # Define constraints
        constraints_list = [
            cp.sum(weights) == 1,  # Weights sum to 1
            weights >= 0,  # No short selling (can be relaxed)
        ]
        
        # Add custom constraints
        if constraints:
            if 'min_weight' in constraints:
                constraints_list.append(weights >= constraints['min_weight'])
            if 'max_weight' in constraints:
                constraints_list.append(weights <= constraints['max_weight'])
            if 'target_return' in constraints:
                constraints_list.append(portfolio_return >= constraints['target_return'])
            if 'max_risk' in constraints:
                max_variance = constraints['max_risk'] ** 2
                constraints_list.append(portfolio_risk <= max_variance)
        
        # Solve optimization
        problem = cp.Problem(objective, constraints_list)
        problem.solve()
        
        if problem.status != 'optimal':
            logger.warning(f"Optimization status: {problem.status}")
            # Return equal weights as fallback
            weights_array = np.ones(n_assets) / n_assets
        else:
            weights_array = weights.value
        
        # Calculate portfolio metrics
        portfolio_return = expected_returns @ weights_array
        portfolio_risk = np.sqrt(weights_array @ cov_matrix @ weights_array)
        sharpe_ratio = (portfolio_return - self.risk_free_rate) / portfolio_risk
        
        return {
            'weights': weights_array,
            'expected_return': portfolio_return * 252,  # Annualized
            'expected_risk': portfolio_risk * np.sqrt(252),  # Annualized
            'sharpe_ratio': sharpe_ratio * np.sqrt(252)
        }
    
    def _minimum_variance_optimization(
        self,
        returns: np.ndarray,
        constraints: Optional[Dict] = None
    ) -> Dict:
        """
        Minimize portfolio variance
        """
        n_assets = returns.shape[1]
        cov_matrix = np.cov(returns.T)
        cov_matrix += np.eye(n_assets) * 1e-6
        
        # Define optimization
        weights = cp.Variable(n_assets)
        portfolio_risk = cp.quad_form(weights, cov_matrix)
        
        objective = cp.Minimize(portfolio_risk)
        
        constraints_list = [
            cp.sum(weights) == 1,
            weights >= 0,
        ]
        
        # Solve
        problem = cp.Problem(objective, constraints_list)
        problem.solve()
        
        weights_array = weights.value if problem.status == 'optimal' else np.ones(n_assets) / n_assets
        
        # Calculate metrics
        expected_returns = np.mean(returns, axis=0)
        portfolio_return = expected_returns @ weights_array
        portfolio_risk = np.sqrt(weights_array @ cov_matrix @ weights_array)
        sharpe_ratio = (portfolio_return - self.risk_free_rate / 252) / portfolio_risk
        
        return {
            'weights': weights_array,
            'expected_return': portfolio_return * 252,
            'expected_risk': portfolio_risk * np.sqrt(252),
            'sharpe_ratio': sharpe_ratio * np.sqrt(252)
        }
    
    def _maximum_sharpe_optimization(
        self,
        returns: np.ndarray,
        constraints: Optional[Dict] = None
    ) -> Dict:
        """
        Maximize Sharpe ratio
        """
        # This is a bit tricky as Sharpe ratio is not convex
        # We use a trick: for each target return, find min variance
        # Then select the portfolio with highest Sharpe
        
        n_assets = returns.shape[1]
        expected_returns = np.mean(returns, axis=0)
        cov_matrix = np.cov(returns.T)
        
        # Generate efficient frontier
        target_returns = np.linspace(
            np.min(expected_returns),
            np.max(expected_returns),
            50
        )
        
        best_sharpe = -np.inf
        best_weights = None
        
        for target_ret in target_returns:
            weights = cp.Variable(n_assets)
            portfolio_risk = cp.quad_form(weights, cov_matrix)
            
            objective = cp.Minimize(portfolio_risk)
            constraints_list = [
                cp.sum(weights) == 1,
                weights >= 0,
                expected_returns @ weights >= target_ret
            ]
            
            problem = cp.Problem(objective, constraints_list)
            problem.solve()
            
            if problem.status == 'optimal':
                w = weights.value
                ret = expected_returns @ w
                risk = np.sqrt(w @ cov_matrix @ w)
                sharpe = (ret - self.risk_free_rate / 252) / risk
                
                if sharpe > best_sharpe:
                    best_sharpe = sharpe
                    best_weights = w
        
        if best_weights is None:
            best_weights = np.ones(n_assets) / n_assets
        
        # Calculate final metrics
        portfolio_return = expected_returns @ best_weights
        portfolio_risk = np.sqrt(best_weights @ cov_matrix @ best_weights)
        
        return {
            'weights': best_weights,
            'expected_return': portfolio_return * 252,
            'expected_risk': portfolio_risk * np.sqrt(252),
            'sharpe_ratio': best_sharpe * np.sqrt(252)
        }
    
    def _risk_parity_optimization(self, returns: np.ndarray) -> Dict:
        """
        Risk parity optimization - equal risk contribution
        """
        n_assets = returns.shape[1]
        cov_matrix = np.cov(returns.T)
        
        # Initial guess - equal weights
        weights = np.ones(n_assets) / n_assets
        
        # Iterative algorithm
        for _ in range(100):
            # Calculate marginal risk contributions
            portfolio_vol = np.sqrt(weights @ cov_matrix @ weights)
            marginal_contrib = cov_matrix @ weights / portfolio_vol
            contrib = weights * marginal_contrib
            
            # Target equal contribution
            target_contrib = portfolio_vol / n_assets
            
            # Update weights
            weights = weights * (target_contrib / contrib)
            weights = weights / np.sum(weights)
        
        # Calculate metrics
        expected_returns = np.mean(returns, axis=0)
        portfolio_return = expected_returns @ weights
        portfolio_risk = np.sqrt(weights @ cov_matrix @ weights)
        sharpe_ratio = (portfolio_return - self.risk_free_rate / 252) / portfolio_risk
        
        return {
            'weights': weights,
            'expected_return': portfolio_return * 252,
            'expected_risk': portfolio_risk * np.sqrt(252),
            'sharpe_ratio': sharpe_ratio * np.sqrt(252)
        }
    
    def _black_litterman_optimization(
        self,
        returns: np.ndarray,
        constraints: Optional[Dict] = None,
        views: Optional[Dict] = None
    ) -> Dict:
        """
        Black-Litterman optimization
        """
        # Simplified implementation
        # In practice, would incorporate market views
        
        n_assets = returns.shape[1]
        
        # Market equilibrium weights (market cap weighted)
        # For demo, use equal weights
        market_weights = np.ones(n_assets) / n_assets
        
        # Calculate implied returns
        cov_matrix = np.cov(returns.T)
        risk_aversion = 2.5
        implied_returns = risk_aversion * cov_matrix @ market_weights
        
        # Without views, this reduces to market weights
        # With views, would blend implied returns with views
        
        if views:
            # Implement view blending
            pass
        
        # For now, return mean-variance with implied returns
        expected_returns = implied_returns
        
        # Run mean-variance with these returns
        weights = cp.Variable(n_assets)
        portfolio_return = expected_returns @ weights
        portfolio_risk = cp.quad_form(weights, cov_matrix)
        
        objective = cp.Maximize(portfolio_return - risk_aversion * portfolio_risk)
        constraints_list = [
            cp.sum(weights) == 1,
            weights >= 0,
        ]
        
        problem = cp.Problem(objective, constraints_list)
        problem.solve()
        
        weights_array = weights.value if problem.status == 'optimal' else market_weights
        
        # Calculate metrics
        portfolio_return = expected_returns @ weights_array
        portfolio_risk = np.sqrt(weights_array @ cov_matrix @ weights_array)
        sharpe_ratio = (portfolio_return - self.risk_free_rate / 252) / portfolio_risk
        
        return {
            'weights': weights_array,
            'expected_return': portfolio_return * 252,
            'expected_risk': portfolio_risk * np.sqrt(252),
            'sharpe_ratio': sharpe_ratio * np.sqrt(252)
        }
    
    def calculate_efficient_frontier(
        self,
        returns: np.ndarray,
        num_portfolios: int = 100
    ) -> List[Dict]:
        """
        Calculate the efficient frontier
        """
        n_assets = returns.shape[1]
        expected_returns = np.mean(returns, axis=0)
        cov_matrix = np.cov(returns.T)
        
        # Range of target returns
        min_ret = np.min(expected_returns)
        max_ret = np.max(expected_returns)
        target_returns = np.linspace(min_ret, max_ret, num_portfolios)
        
        frontier = []
        
        for target_ret in target_returns:
            weights = cp.Variable(n_assets)
            portfolio_risk = cp.quad_form(weights, cov_matrix)
            
            objective = cp.Minimize(portfolio_risk)
            constraints_list = [
                cp.sum(weights) == 1,
                weights >= 0,
                expected_returns @ weights >= target_ret
            ]
            
            problem = cp.Problem(objective, constraints_list)
            problem.solve()
            
            if problem.status == 'optimal':
                w = weights.value
                risk = np.sqrt(w @ cov_matrix @ w)
                
                frontier.append({
                    'return': target_ret * 252,
                    'risk': risk * np.sqrt(252),
                    'weights': w.tolist()
                })
        
        return frontier