"""
Visualización de resultados de optimización de portafolio
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Optional
import warnings
warnings.filterwarnings('ignore')

# Configuración de estilo
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = (12, 8)
plt.rcParams['font.size'] = 10


def plot_portfolio_weights(
    weights: np.ndarray,
    tickers: List[str],
    title: str = "Portfolio Weights",
    save_path: Optional[str] = None
):
    """
    Gráfico de barras con pesos del portafolio
    
    Args:
        weights: Array con pesos
        tickers: Lista de nombres de activos
        title: Título del gráfico
        save_path: Ruta para guardar (opcional)
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    
    colors = sns.color_palette("husl", len(weights))
    bars = ax.bar(tickers, weights * 100, color=colors, edgecolor='black')
    
    # Etiquetas
    for bar in bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height,
                f'{height:.1f}%',
                ha='center', va='bottom', fontweight='bold')
    
    ax.set_ylabel('Weight (%)', fontsize=12, fontweight='bold')
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.set_ylim(0, max(weights * 100) * 1.15)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Gráfico guardado: {save_path}")
    
    plt.show()


def plot_efficient_frontier(
    expected_returns: np.ndarray,
    cov_matrix: np.ndarray,
    optimal_weights: np.ndarray,
    n_portfolios: int = 1000,
    save_path: Optional[str] = None
):
    """
    Grafica la frontera eficiente con el portafolio óptimo marcado
    
    Args:
        expected_returns: Retornos esperados
        cov_matrix: Matriz de covarianza
        optimal_weights: Pesos del portafolio óptimo
        n_portfolios: Número de portafolios aleatorios a simular
        save_path: Ruta para guardar
    """
    n_assets = len(expected_returns)
    
    # Generar portafolios aleatorios
    np.random.seed(42)
    results = np.zeros((3, n_portfolios))
    
    for i in range(n_portfolios):
        # Pesos aleatorios
        weights = np.random.random(n_assets)
        weights /= np.sum(weights)
        
        # Calcular retorno y riesgo
        portfolio_return = np.dot(weights, expected_returns)
        portfolio_std = np.sqrt(np.dot(weights.T, np.dot(cov_matrix, weights)))
        
        results[0, i] = portfolio_return
        results[1, i] = portfolio_std
        results[2, i] = portfolio_return / portfolio_std  # Sharpe
    
    # Portafolio óptimo
    opt_return = np.dot(optimal_weights, expected_returns)
    opt_std = np.sqrt(np.dot(optimal_weights.T, np.dot(cov_matrix, optimal_weights)))
    
    # Graficar
    fig, ax = plt.subplots(figsize=(12, 8))
    
    scatter = ax.scatter(
        results[1, :],
        results[0, :],
        c=results[2, :],
        cmap='viridis',
        alpha=0.5,
        s=30
    )
    
    # Marcar óptimo
    ax.scatter(
        opt_std,
        opt_return,
        marker='*',
        color='red',
        s=500,
        edgecolors='black',
        linewidths=2,
        label='Optimal Portfolio',
        zorder=5
    )
    
    # Colorbar
    cbar = plt.colorbar(scatter, ax=ax)
    cbar.set_label('Sharpe Ratio', fontsize=12, fontweight='bold')
    
    ax.set_xlabel('Risk (Volatility)', fontsize=12, fontweight='bold')
    ax.set_ylabel('Expected Return', fontsize=12, fontweight='bold')
    ax.set_title('Efficient Frontier', fontsize=14, fontweight='bold')
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Frontera eficiente guardada: {save_path}")
    
    plt.show()


def plot_portfolio_evolution(
    weights: np.ndarray,
    returns_data: pd.DataFrame,
    tickers: List[str],
    title: str = "Portfolio Value Evolution",
    save_path: Optional[str] = None
):
    """
    Grafica la evolución del valor del portafolio
    
    Args:
        weights: Pesos del portafolio
        returns_data: DataFrame con retornos históricos
        tickers: Lista de nombres
        title: Título
        save_path: Ruta para guardar
    """
    # Calcular valor del portafolio
    portfolio_returns = returns_data.dot(weights)
    portfolio_value = (1 + portfolio_returns).cumprod()
    
    # También graficar activos individuales
    individual_values = (1 + returns_data).cumprod()
    
    fig, ax = plt.subplots(figsize=(14, 7))
    
    # Activos individuales (líneas finas)
    for col in individual_values.columns:
        ax.plot(individual_values.index, individual_values[col], 
                alpha=0.3, linewidth=1, label=col)
    
    # Portafolio (línea gruesa)
    ax.plot(portfolio_value.index, portfolio_value, 
            color='red', linewidth=3, label='Portfolio', zorder=5)
    
    ax.set_xlabel('Date', fontsize=12, fontweight='bold')
    ax.set_ylabel('Portfolio Value', fontsize=12, fontweight='bold')
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.legend(loc='best', fontsize=9)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Evolución guardada: {save_path}")
    
    plt.show()


def plot_correlation_heatmap(
    returns_data: pd.DataFrame,
    save_path: Optional[str] = None
):
    """
    Mapa de calor de correlaciones
    
    Args:
        returns_data: DataFrame con retornos
        save_path: Ruta para guardar
    """
    corr_matrix = returns_data.corr()
    
    fig, ax = plt.subplots(figsize=(10, 8))
    
    sns.heatmap(
        corr_matrix,
        annot=True,
        fmt='.2f',
        cmap='coolwarm',
        center=0,
        square=True,
        linewidths=1,
        cbar_kws={"shrink": 0.8},
        ax=ax
    )
    
    ax.set_title('Asset Correlation Matrix', fontsize=14, fontweight='bold')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Heatmap guardado: {save_path}")
    
    plt.show()


def plot_comparison(
    results: Dict[str, Dict],
    metric: str = 'sharpe_ratio',
    save_path: Optional[str] = None
):
    """
    Compara métrica entre diferentes métodos
    
    Args:
        results: Dict {método: {metrics}}
        metric: Métrica a comparar
        save_path: Ruta para guardar
    """
    methods = list(results.keys())
    values = [results[m][metric] for m in methods]
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    colors = sns.color_palette("Set2", len(methods))
    bars = ax.barh(methods, values, color=colors, edgecolor='black')
    
    # Etiquetas
    for bar, val in zip(bars, values):
        width = bar.get_width()
        ax.text(width, bar.get_y() + bar.get_height()/2.,
                f'{val:.3f}',
                ha='left', va='center', fontweight='bold', fontsize=11)
    
    ax.set_xlabel(metric.replace('_', ' ').title(), fontsize=12, fontweight='bold')
    ax.set_title(f'Comparison: {metric.replace("_", " ").title()}', 
                 fontsize=14, fontweight='bold')
    ax.grid(axis='x', alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Comparación guardada: {save_path}")
    
    plt.show()


if __name__ == '__main__':
    # Demo
    print("=" * 60)
    print("VISUALIZATION DEMO")
    print("=" * 60)
    
    # Datos sintéticos
    np.random.seed(42)
    tickers = ['AAPL', 'MSFT', 'GOOGL', 'AMZN']
    weights = np.array([0.30, 0.25, 0.25, 0.20])
    
    # Generar retornos
    n_days = 252
    returns_data = pd.DataFrame(
        np.random.randn(n_days, 4) * 0.02 + 0.0003,
        columns=tickers
    )
    
    expected_returns = returns_data.mean() * 252
    cov_matrix = returns_data.cov() * 252
    
    # Demos
    print("\n1. Portfolio Weights")
    plot_portfolio_weights(weights, tickers)
    
    print("\n2. Efficient Frontier")
    plot_efficient_frontier(expected_returns.values, cov_matrix.values, weights)
    
    print("\n3. Correlation Heatmap")
    plot_correlation_heatmap(returns_data)
    
    print(f"\n✓ Visualizaciones funcionando!")
