"""
Optimizador de Portafolio usando QAOA (Quantum Approximate Optimization Algorithm)
Combina predicciones clásicas con optimización cuántica
"""

import numpy as np
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
from qiskit_aer import AerSimulator
from qiskit.primitives import StatevectorSampler
from qiskit_algorithms.optimizers import COBYLA, SLSQP
from typing import Dict, List, Tuple
import warnings
warnings.filterwarnings('ignore')


class PortfolioQAOA:
    """
    Optimizador cuántico de portafolio usando QAOA
    
    Resuelve el problema:
        maximize: expected_return - λ * risk
        subject to: sum(weights) = 1, 0 <= weights <= max_weight
    """
    
    def __init__(
        self,
        n_assets: int,
        risk_aversion: float = 0.5,
        max_weight: float = 0.30,
        min_weight: float = 0.05,
        n_bits_per_asset: int = 4,
        qaoa_reps: int = 3,
        use_simulator: bool = True
    ):
        """
        Args:
            n_assets: Número de activos en el portafolio
            risk_aversion: Parámetro λ de aversión al riesgo (0=solo retorno, 1=solo riesgo)
            max_weight: Peso máximo por activo (default: 30%)
            min_weight: Peso mínimo por activo si se incluye (default: 5%)
            n_bits_per_asset: Bits para codificar cada peso (más bits = más precisión)
            qaoa_reps: Número de repeticiones del circuito QAOA (p layers)
            use_simulator: Si usar simulador o hardware cuántico
        """
        self.n_assets = n_assets
        self.risk_aversion = risk_aversion
        self.max_weight = max_weight
        self.min_weight = min_weight
        self.n_bits_per_asset = n_bits_per_asset
        self.qaoa_reps = qaoa_reps
        
        # Total de qubits necesarios
        self.n_qubits = n_assets * n_bits_per_asset
        
        # Backend
        if use_simulator:
            self.backend = AerSimulator()
        else:
            # Aquí se podría usar un backend real de IBM Quantum
            # from qiskit_ibm_provider import IBMProvider
            # provider = IBMProvider()
            # self.backend = provider.get_backend('ibm_brisbane')
            self.backend = AerSimulator()
        
        print(f"✓ QAOA Optimizer inicializado:")
        print(f"  Activos: {n_assets}")
        print(f"  Qubits totales: {self.n_qubits}")
        print(f"  Precisión: {n_bits_per_asset} bits/activo")
        print(f"  Backend: {self.backend}")
    
    def encode_weights(self, bitstring: str) -> np.ndarray:
        """
        Decodifica bitstring a pesos de portafolio
        
        Args:
            bitstring: String binario de longitud n_qubits
            
        Returns:
            Array de pesos normalizados
        """
        weights = []
        
        # Dividir bitstring en chunks para cada activo
        for i in range(self.n_assets):
            start = i * self.n_bits_per_asset
            end = start + self.n_bits_per_asset
            bits = bitstring[start:end]
            
            # Convertir binario a decimal
            decimal = int(bits, 2)
            
            # Normalizar al rango [0, max_weight]
            max_val = 2 ** self.n_bits_per_asset - 1
            weight = (decimal / max_val) * self.max_weight
            
            weights.append(weight)
        
        weights = np.array(weights)
        
        # Normalizar para que sumen 1
        total = weights.sum()
        if total > 0:
            weights = weights / total
        else:
            # Si todos son 0, distribuir uniformemente
            weights = np.ones(self.n_assets) / self.n_assets
        
        return weights
    
    def calculate_portfolio_objective(
        self,
        weights: np.ndarray,
        expected_returns: np.ndarray,
        cov_matrix: np.ndarray
    ) -> float:
        """
        Calcula función objetivo del portafolio (negativa para minimizar)
        
        Args:
            weights: Pesos del portafolio
            expected_returns: Retornos esperados de cada activo
            cov_matrix: Matriz de covarianza
            
        Returns:
            -1 * (retorno_esperado - λ * riesgo)
        """
        # Retorno esperado
        portfolio_return = np.dot(weights, expected_returns)
        
        # Riesgo (varianza del portafolio)
        portfolio_variance = np.dot(weights.T, np.dot(cov_matrix, weights))
        
        # Función objetivo (maximizar retorno - λ*riesgo)
        # Devolvemos negativo porque QAOA minimiza
        objective = -(portfolio_return - self.risk_aversion * portfolio_variance)
        
        return objective
    
    def apply_constraints_penalty(
        self,
        weights: np.ndarray,
        penalty_strength: float = 1000.0
    ) -> float:
        """
        Calcula penalización por violación de restricciones
        
        Args:
            weights: Pesos del portafolio
            penalty_strength: Fuerza de la penalización
            
        Returns:
            Valor de penalización
        """
        penalty = 0.0
        
        # Penalizar si suma != 1
        sum_penalty = abs(weights.sum() - 1.0)
        penalty += penalty_strength * sum_penalty ** 2
        
        # Penalizar pesos fuera de rango
        for w in weights:
            if w < 0:
                penalty += penalty_strength * w ** 2
            if w > self.max_weight:
                penalty += penalty_strength * (w - self.max_weight) ** 2
        
        # Penalizar diversificación insuficiente
        n_active_assets = np.sum(weights > self.min_weight)
        if n_active_assets < 3:  # Mínimo 3 activos
            penalty += penalty_strength * (3 - n_active_assets) ** 2
        
        return penalty
    
    def cost_function(
        self,
        bitstring: str,
        expected_returns: np.ndarray,
        cov_matrix: np.ndarray
    ) -> float:
        """
        Función de costo completa para QAOA
        
        Args:
            bitstring: Configuración de qubits
            expected_returns: Retornos esperados
            cov_matrix: Matriz de covarianza
            
        Returns:
            Valor de costo (objetivo + penalizaciones)
        """
        # Decodificar pesos
        weights = self.encode_weights(bitstring)
        
        # Objetivo principal
        objective = self.calculate_portfolio_objective(
            weights, expected_returns, cov_matrix
        )
        
        # Penalizaciones
        penalty = self.apply_constraints_penalty(weights)
        
        return objective + penalty
    
    def optimize(
        self,
        expected_returns: np.ndarray,
        cov_matrix: np.ndarray,
        max_iterations: int = 1000,
        optimizer_name: str = 'COBYLA'
    ) -> Dict:
        """
        Ejecuta optimización QAOA
        
        Args:
            expected_returns: Array con retornos esperados (de modelo clásico)
            cov_matrix: Matriz de covarianza
            max_iterations: Iteraciones máximas del optimizador clásico
            optimizer_name: 'COBYLA' o 'SLSQP'
            
        Returns:
            Dict con resultados de optimización
        """
        print(f"\n{'='*60}")
        print("EJECUTANDO OPTIMIZACIÓN QAOA")
        print(f"{'='*60}")
        
        # Validar inputs
        assert len(expected_returns) == self.n_assets
        assert cov_matrix.shape == (self.n_assets, self.n_assets)
        
        # Crear función de costo parametrizada
        def qaoa_objective(counts: Dict[str, int]) -> float:
            """Función objetivo para QAOA basada en counts"""
            total_cost = 0.0
            total_counts = sum(counts.values())
            
            for bitstring, count in counts.items():
                cost = self.cost_function(bitstring, expected_returns, cov_matrix)
                total_cost += cost * (count / total_counts)
            
            return total_cost
        
        # Crear circuito QAOA
        print(f"Creando circuito QAOA con {self.qaoa_reps} repeticiones...")
        qc = self._create_qaoa_circuit(expected_returns, cov_matrix)
        
        # Configurar optimizador clásico
        if optimizer_name == 'COBYLA':
            classical_optimizer = COBYLA(maxiter=max_iterations)
        elif optimizer_name == 'SLSQP':
            classical_optimizer = SLSQP(maxiter=max_iterations)
        else:
            raise ValueError(f"Optimizador no reconocido: {optimizer_name}")
        
        # Ejecutar QAOA
        print(f"Optimizando con {optimizer_name}...")
        
        # Usar enfoque directo con sampling
        best_cost = float('inf')
        best_weights = None
        best_bitstring = None
        
        # Parámetros iniciales aleatorios
        n_params = 2 * self.qaoa_reps  # beta y gamma para cada capa
        
        def objective_function(params):
            nonlocal best_cost, best_weights, best_bitstring
            
            # Asignar parámetros al circuito
            bound_qc = qc.assign_parameters(params)
            
            # Simular con StatevectorSampler (Qiskit 2.x)
            bound_qc.measure_all()
            sampler = StatevectorSampler()
            job = sampler.run([bound_qc], shots=1024)
            result = job.result()
            
            # Obtener counts - Qiskit 2.x API
            pub_result = result[0]
            counts_dict = pub_result.data.meas.get_counts()
            
            # Convertir a formato binario string
            counts = {format(int(state), f'0{self.n_qubits}b'): count 
                     for state, count in counts_dict.items()}
            
            # Calcular costo
            cost = qaoa_objective(counts)
            
            # Actualizar mejor solución
            if cost < best_cost:
                best_cost = cost
                # Encontrar bitstring más frecuente
                best_bitstring = max(counts, key=counts.get)
                best_weights = self.encode_weights(best_bitstring)
                
                print(f"  Mejor costo encontrado: {cost:.6f}")
            
            return cost
        
        # Optimizar
        initial_params = np.random.uniform(0, 2*np.pi, n_params)
        
        if optimizer_name == 'COBYLA':
            from scipy.optimize import minimize
            opt_result = minimize(
                objective_function,
                initial_params,
                method='COBYLA',
                options={'maxiter': max_iterations}
            )
        else:
            opt_result = classical_optimizer.minimize(
                objective_function,
                initial_params
            )
        
        # Resultados
        print(f"\n{'='*60}")
        print("RESULTADOS DE OPTIMIZACIÓN")
        print(f"{'='*60}")
        print(f"Costo final: {best_cost:.6f}")
        print(f"Bitstring óptimo: {best_bitstring}")
        print(f"\nPesos optimizados:")
        for i, w in enumerate(best_weights):
            print(f"  Activo {i+1}: {w*100:.2f}%")
        
        # Calcular métricas del portafolio
        portfolio_return = np.dot(best_weights, expected_returns)
        portfolio_risk = np.sqrt(np.dot(best_weights.T, np.dot(cov_matrix, best_weights)))
        sharpe_ratio = portfolio_return / portfolio_risk if portfolio_risk > 0 else 0
        
        print(f"\nMétricas del Portafolio:")
        print(f"  Retorno Esperado: {portfolio_return*100:.2f}%")
        print(f"  Riesgo (Volatilidad): {portfolio_risk*100:.2f}%")
        print(f"  Sharpe Ratio: {sharpe_ratio:.2f}")
        
        return {
            'weights': best_weights,
            'cost': best_cost,
            'bitstring': best_bitstring,
            'expected_return': portfolio_return,
            'risk': portfolio_risk,
            'sharpe_ratio': sharpe_ratio,
            'optimization_result': opt_result
        }
    
    def _create_qaoa_circuit(
        self,
        expected_returns: np.ndarray,
        cov_matrix: np.ndarray
    ) -> QuantumCircuit:
        """
        Crea circuito QAOA parametrizado
        
        Returns:
            QuantumCircuit con parámetros
        """
        qc = QuantumCircuit(self.n_qubits)
        
        # Estado inicial: superposición uniforme
        qc.h(range(self.n_qubits))
        
        # Parámetros
        betas = [Parameter(f'β_{i}') for i in range(self.qaoa_reps)]
        gammas = [Parameter(f'γ_{i}') for i in range(self.qaoa_reps)]
        
        for p in range(self.qaoa_reps):
            # Cost Hamiltonian (simplified)
            # Codifica la función de costo en rotaciones
            for i in range(self.n_qubits):
                qc.rz(2 * gammas[p], i)
            
            # Entanglement entre qubits relacionados
            for i in range(self.n_assets - 1):
                q1 = i * self.n_bits_per_asset
                q2 = (i + 1) * self.n_bits_per_asset
                if q1 < self.n_qubits and q2 < self.n_qubits:
                    qc.cx(q1, q2)
            
            # Mixer Hamiltonian
            for i in range(self.n_qubits):
                qc.rx(2 * betas[p], i)
        
        return qc


if __name__ == '__main__':
    # Demo
    print("=" * 60)
    print("QAOA PORTFOLIO OPTIMIZER DEMO")
    print("=" * 60)
    
    # Datos sintéticos
    n_assets = 4  # Reducido para demo rápida
    
    # Retornos esperados aleatorios
    expected_returns = np.random.uniform(0.05, 0.20, n_assets)
    
    # Matriz de covarianza aleatoria simétrica positiva definida
    A = np.random.rand(n_assets, n_assets)
    cov_matrix = (A + A.T) / 2 + np.eye(n_assets) * 0.1
    
    print(f"\nRetornos esperados: {expected_returns}")
    print(f"Matriz de covarianza:\n{cov_matrix}")
    
    # Crear optimizador
    optimizer = PortfolioQAOA(
        n_assets=n_assets,
        risk_aversion=0.5,
        n_bits_per_asset=3,  # Reducido para demo
        qaoa_reps=2
    )
    
    # Optimizar
    result = optimizer.optimize(expected_returns, cov_matrix, max_iterations=50)
    
    print(f"\n✓ QAOA funcionando correctamente!")
