"""
Data Loader para Quantum Portfolio Optimizer
Descarga y procesa datos financieros históricos desde Yahoo Finance
"""

import yfinance as yf
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import List, Tuple, Dict
import os


class FinancialDataLoader:
    """Descarga y prepara datos financieros para optimización de portafolio"""
    
    def __init__(self, tickers: List[str], start_date: str = None, end_date: str = None):
        """
        Args:
            tickers: Lista de símbolos de acciones (ej: ['AAPL', 'MSFT', 'GOOGL'])
            start_date: Fecha inicio en formato 'YYYY-MM-DD' (default: 5 años atrás)
            end_date: Fecha fin en formato 'YYYY-MM-DD' (default: hoy)
        """
        self.tickers = tickers
        self.end_date = end_date or datetime.now().strftime('%Y-%m-%d')
        
        if start_date is None:
            start = datetime.now() - timedelta(days=252*5)  # 5 años
            self.start_date = start.strftime('%Y-%m-%d')
        else:
            self.start_date = start_date
        
        self.raw_data = None
        self.returns = None
        self.log_returns = None
        
    def download_data(self, save_path: str = None) -> pd.DataFrame:
        """
        Descarga precios históricos ajustados
        
        Args:
            save_path: Ruta opcional para guardar datos en CSV
            
        Returns:
            DataFrame con precios ajustados para cada ticker
        """
        print(f"Descargando datos para {len(self.tickers)} activos...")
        print(f"Periodo: {self.start_date} a {self.end_date}")
        
        # Descargar datos
        data = yf.download(
            self.tickers,
            start=self.start_date,
            end=self.end_date,
            progress=True
        )
        
        # Usar precios de cierre ajustados
        if len(self.tickers) == 1:
            self.raw_data = pd.DataFrame({self.tickers[0]: data['Adj Close']})
        else:
            self.raw_data = data['Adj Close']
        
        # Limpiar valores faltantes (forward fill luego backward fill)
        self.raw_data = self.raw_data.fillna(method='ffill').fillna(method='bfill')
        
        print(f"✓ Datos descargados: {len(self.raw_data)} días")
        
        # Guardar si se especifica ruta
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            self.raw_data.to_csv(save_path)
            print(f"✓ Datos guardados en: {save_path}")
        
        return self.raw_data
    
    def calculate_returns(self, log_returns: bool = True) -> pd.DataFrame:
        """
        Calcula retornos de los activos
        
        Args:
            log_returns: Si True, calcula retornos logarítmicos; si False, retornos simples
            
        Returns:
            DataFrame con retornos
        """
        if self.raw_data is None:
            raise ValueError("Primero debes descargar los datos con download_data()")
        
        if log_returns:
            self.log_returns = np.log(self.raw_data / self.raw_data.shift(1))
            self.returns = self.log_returns.dropna()
        else:
            self.returns = self.raw_data.pct_change().dropna()
        
        print(f"✓ Retornos calculados: {len(self.returns)} observaciones")
        return self.returns
    
    def get_covariance_matrix(self, annualize: bool = True) -> pd.DataFrame:
        """
        Calcula matriz de covarianza de los retornos
        
        Args:
            annualize: Si True, anualiza la covarianza (multiplicar por 252)
            
        Returns:
            Matriz de covarianza
        """
        if self.returns is None:
            self.calculate_returns()
        
        cov_matrix = self.returns.cov()
        
        if annualize:
            cov_matrix = cov_matrix * 252  # 252 días de trading
        
        return cov_matrix
    
    def get_expected_returns(self, method: str = 'mean', annualize: bool = True) -> pd.Series:
        """
        Calcula retornos esperados
        
        Args:
            method: 'mean' para promedio histórico, 'ewm' para media exponencial
            annualize: Si True, anualiza los retornos
            
        Returns:
            Series con retornos esperados por activo
        """
        if self.returns is None:
            self.calculate_returns()
        
        if method == 'mean':
            expected = self.returns.mean()
        elif method == 'ewm':
            # Media móvil exponencial con decay de 60 días
            expected = self.returns.ewm(span=60).mean().iloc[-1]
        else:
            raise ValueError(f"Método no reconocido: {method}")
        
        if annualize:
            expected = expected * 252
        
        return expected
    
    def get_correlation_matrix(self) -> pd.DataFrame:
        """Calcula matriz de correlación"""
        if self.returns is None:
            self.calculate_returns()
        return self.returns.corr()
    
    def get_data_summary(self) -> Dict:
        """
        Genera resumen estadístico de los datos
        
        Returns:
            Diccionario con estadísticas
        """
        if self.returns is None:
            self.calculate_returns()
        
        summary = {
            'tickers': self.tickers,
            'start_date': self.start_date,
            'end_date': self.end_date,
            'n_observations': len(self.returns),
            'expected_returns_annual': self.get_expected_returns(annualize=True).to_dict(),
            'volatility_annual': (self.returns.std() * np.sqrt(252)).to_dict(),
            'sharpe_ratio': (self.get_expected_returns(annualize=True) / 
                           (self.returns.std() * np.sqrt(252))).to_dict()
        }
        
        return summary
    
    def train_test_split(self, test_size: float = 0.2) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Divide datos en entrenamiento y prueba (temporal)
        
        Args:
            test_size: Proporción de datos para test (default: 20%)
            
        Returns:
            Tupla (train_data, test_data)
        """
        if self.returns is None:
            self.calculate_returns()
        
        split_idx = int(len(self.returns) * (1 - test_size))
        train = self.returns.iloc[:split_idx]
        test = self.returns.iloc[split_idx:]
        
        print(f"✓ Train: {len(train)} días, Test: {len(test)} días")
        return train, test


def load_sample_data(tickers: List[str] = None, save_dir: str = None) -> FinancialDataLoader:
    """
    Función helper para cargar datos de ejemplo rápidamente
    
    Args:
        tickers: Lista de tickers (default: FAANG + Microsoft + JPM + J&J)
        save_dir: Directorio para guardar datos
        
    Returns:
        FinancialDataLoader con datos cargados
    """
    if tickers is None:
        tickers = ['AAPL', 'MSFT', 'GOOGL', 'AMZN', 'META', 'JPM', 'JNJ', 'TSLA']
    
    loader = FinancialDataLoader(tickers)
    
    # Descargar datos
    save_path = None
    if save_dir:
        save_path = os.path.join(save_dir, 'historical_prices.csv')
    
    loader.download_data(save_path=save_path)
    loader.calculate_returns()
    
    return loader


if __name__ == '__main__':
    # Ejemplo de uso
    print("=" * 60)
    print("QUANTUM PORTFOLIO OPTIMIZER - Data Loader Demo")
    print("=" * 60)
    
    # Cargar datos de ejemplo
    loader = load_sample_data(save_dir='data/raw')
    
    # Mostrar resumen
    print("\n" + "=" * 60)
    print("RESUMEN DE DATOS")
    print("=" * 60)
    
    summary = loader.get_data_summary()
    print(f"\nActivos: {', '.join(summary['tickers'])}")
    print(f"Periodo: {summary['start_date']} → {summary['end_date']}")
    print(f"Observaciones: {summary['n_observations']} días")
    
    print("\nRetornos Anualizados:")
    for ticker, ret in summary['expected_returns_annual'].items():
        print(f"  {ticker}: {ret*100:.2f}%")
    
    print("\nVolatilidad Anualizada:")
    for ticker, vol in summary['volatility_annual'].items():
        print(f"  {ticker}: {vol*100:.2f}%")
    
    print("\nSharpe Ratio:")
    for ticker, sr in summary['sharpe_ratio'].items():
        print(f"  {ticker}: {sr:.2f}")
    
    # Matriz de correlación
    print("\n" + "=" * 60)
    print("MATRIZ DE CORRELACIÓN")
    print("=" * 60)
    print(loader.get_correlation_matrix().round(2))
    
    print("\n✓ Data loader funcionando correctamente!")
