"""
Main HybridOptimizer class - Simplified API for portfolio optimization
"""

import numpy as np
import pandas as pd
from typing import List, Optional, Dict, Union
import warnings

from ..data.loader import FinancialDataLoader as DataLoader
from ..models.quantum import PortfolioQAOA as QAOAOptimizer
from ..baselines.classical import (
    MarkowitzOptimizer,
    EqualWeightOptimizer,
    MinVarianceOptimizer,
    MaxSharpeOptimizer
)


class HybridOptimizer:
    """
    Main API for portfolio optimization
    
    Combines classical and quantum optimization methods with simple interface.
    
    Args:
        tickers: List of stock symbols (e.g. ['AAPL', 'MSFT', 'GOOGL'])
        method: Optimization method - 'quantum', 'classical', or 'hybrid'
        classical_model: 'transformer' or 'lstm' (for hybrid method)
        use_quantum: Whether to use QAOA (for hybrid method)
        risk_aversion: Risk aversion parameter 0-1 (0=aggressive, 1=conservative)
        max_weight: Maximum weight per asset (default: 0.30 = 30%)
        min_weight: Minimum weight per asset (default: 0.05 = 5%)
        
    Example:
        >>> import quantfolio as qf
        >>> optimizer = qf.HybridOptimizer(['AAPL', 'MSFT', 'GOOGL'])
        >>> weights = optimizer.optimize(risk_aversion=0.5)
        >>> print(weights)
    """
    
    def __init__(
        self,
        tickers: List[str],
        method: str = 'quantum',
        classical_model: str = 'lstm',
        use_quantum: bool = None,
        risk_aversion: float = 0.5,
        max_weight: float = 0.30,
        min_weight: float = 0.05,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
    ):
        self.tickers = tickers
        self.n_assets = len(tickers)
        self.method = method.lower()
        self.classical_model = classical_model
        self.risk_aversion = risk_aversion
        self.max_weight = max_weight
        self.min_weight = min_weight
        
        # Auto-determine use_quantum from method if not specified
        if use_quantum is None:
            use_quantum = (method in ['quantum', 'hybrid'])
        self.use_quantum = use_quantum
        
        # Load data
        self.data_loader = DataLoader(tickers, start_date=start_date, end_date=end_date)
        self._data_loaded = False
        
    def _ensure_data_loaded(self):
        """Lazy load data"""
        if not self._data_loaded:
            self.data_loader.download_data()
            self.data_loader.calculate_returns()
            self._data_loaded = True
    
    def optimize(
        self,
        train_predictor: bool = False,
        epochs: int = 30,
        risk_aversion: Optional[float] = None,
        **kwargs
    ) -> Union[np.ndarray, Dict]:
        """
        Run portfolio optimization
        
        Args:
            train_predictor: Whether to train ML predictor (only for hybrid)
            epochs: Training epochs for predictor
            risk_aversion: Override constructor risk aversion
            **kwargs: Additional arguments for optimizer
            
        Returns:
            If simple call: np.ndarray of weights
            If detailed: Dict with weights, return, risk, sharpe_ratio
        """
        self._ensure_data_loaded()
        
        # Override risk aversion if provided
        if risk_aversion is not None:
            self.risk_aversion = risk_aversion
        
        # Get expected returns and covariance
        expected_returns = self.data_loader.get_expected_returns().values
        cov_matrix = self.data_loader.get_covariance_matrix().values
        
        # Choose optimization method
        if self.method == 'quantum' or (self.method == 'hybrid' and self.use_quantum):
            # Quantum optimization
            optimizer = QAOAOptimizer(
                n_assets=self.n_assets,
                risk_aversion=self.risk_aversion,
                max_weight=self.max_weight,
                min_weight=self.min_weight,
                n_bits_per_asset=4,
                qaoa_reps=2,  # Reduced for speed
                use_simulator=True
            )
            
            result = optimizer.optimize(
                expected_returns,
                cov_matrix,
                max_iterations=kwargs.get('max_iterations', 50)
            )
            
            return result
            
        elif self.method == 'classical':
            # Classical Markowitz optimization
            optimizer = MarkowitzOptimizer(
                expected_returns,
                cov_matrix,
                self.risk_aversion
            )
            
            result = optimizer.optimize(
                max_weight=self.max_weight,
                min_weight=self.min_weight
            )
            
            return result
        
        else:
            raise ValueError(f"Unknown method: {self.method}. Use 'quantum', 'classical', or 'hybrid'")
    
    def get_weights(self) -> np.ndarray:
        """Get just the weights array from optimization"""
        result = self.optimize()
        if isinstance(result, dict):
            return result['weights']
        return result
    
    def plot_weights(self):
        """Visualize portfolio weights"""
        from ..viz.plots import plot_portfolio_weights
        
        weights = self.get_weights()
        plot_portfolio_weights(weights, self.tickers, title=f"{self.method.title()} Portfolio Weights")
    
    def plot_efficient_frontier(self):
        """Plot efficient frontier"""
        from ..viz.plots import plot_efficient_frontier
        
        self._ensure_data_loaded()
        expected_returns = self.data_loader.get_expected_returns().values
        cov_matrix = self.data_loader.get_covariance_matrix().values
        weights = self.get_weights()
        
        plot_efficient_frontier(expected_returns, cov_matrix, weights)
    
    def plot_evolution(self):
        """Plot portfolio value evolution"""
        from ..viz.plots import plot_portfolio_evolution
        
        self._ensure_data_loaded()
        weights = self.get_weights()
        
        plot_portfolio_evolution(
            weights,
            self.data_loader.returns,
            self.tickers
        )


# Convenience function
def compare_methods(
    tickers: List[str],
    methods: List[str] = ['quantum', 'classical'],
    risk_aversion: float = 0.5
) -> pd.DataFrame:
    """
    Compare different optimization methods
    
    Args:
        tickers: List of stock symbols
        methods: List of methods to compare
        risk_aversion: Risk aversion parameter
        
    Returns:
        DataFrame with comparison results
    """
    from ..metrics.portfolio import calculate_portfolio_metrics
    
    # Load data once
    loader = DataLoader(tickers)
    loader.download_data()
    loader.calculate_returns()
    
    results = {}
    
    for method in methods:
        optimizer = HybridOptimizer(
            tickers,
            method=method,
            risk_aversion=risk_aversion
        )
        optimizer.data_loader = loader  # Reuse data
        optimizer._data_loaded = True
        
        result = optimizer.optimize()
        weights = result['weights'] if isinstance(result, dict) else result
        
        # Calculate metrics
        metrics = calculate_portfolio_metrics(weights, loader.returns)
        metrics['method'] = method
        results[method] = metrics
    
    return pd.DataFrame(results).T


__all__ = ['HybridOptimizer', 'compare_methods']
