"""
Stock Market Time Series Dataset - Data Loading Utilities

Author: Molla Samser
Organization: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
"""

import pandas as pd
import numpy as np
import json
from pathlib import Path
from typing import Dict, List, Optional, Union
import warnings
warnings.filterwarnings('ignore')


class StockDataLoader:
    """
    Utility class for loading and managing stock market time series data.
    """
    
    def __init__(self, data_dir: str = 'data', metadata_dir: str = 'metadata'):
        """
        Initialize the StockDataLoader.
        
        Args:
            data_dir: Directory containing stock CSV files
            metadata_dir: Directory containing metadata JSON files
        """
        self.data_dir = Path(data_dir)
        self.metadata_dir = Path(metadata_dir)
        self.stocks_cache = {}
        
    def load_stock(self, symbol: str, parse_dates: bool = True) -> pd.DataFrame:
        """
        Load stock data for a given symbol.
        
        Args:
            symbol: Stock symbol (e.g., 'AAPL', 'GOOGL')
            parse_dates: Whether to parse dates and set as index
            
        Returns:
            DataFrame with stock data
        """
        # Check cache first
        if symbol in self.stocks_cache:
            return self.stocks_cache[symbol].copy()
        
        file_path = self.data_dir / f"{symbol}.csv"
        
        if not file_path.exists():
            raise FileNotFoundError(f"Stock data file not found: {file_path}")
        
        df = pd.read_csv(file_path)
        
        if parse_dates:
            df['Date'] = pd.to_datetime(df['Date'])
            df.set_index('Date', inplace=True)
            df.sort_index(inplace=True)
        
        # Cache the data
        self.stocks_cache[symbol] = df.copy()
        
        return df
    
    def load_multiple_stocks(self, symbols: List[str]) -> Dict[str, pd.DataFrame]:
        """
        Load multiple stock datasets.
        
        Args:
            symbols: List of stock symbols
            
        Returns:
            Dictionary with symbol as key and DataFrame as value
        """
        return {symbol: self.load_stock(symbol) for symbol in symbols}
    
    def load_all_stocks(self) -> Dict[str, pd.DataFrame]:
        """
        Load all available stock datasets.
        
        Returns:
            Dictionary with all stock data
        """
        csv_files = list(self.data_dir.glob("*.csv"))
        symbols = [f.stem for f in csv_files]
        return self.load_multiple_stocks(symbols)
    
    def load_metadata(self, metadata_file: str = 'stock_info.json') -> Dict:
        """
        Load metadata information.
        
        Args:
            metadata_file: Name of metadata JSON file
            
        Returns:
            Dictionary with metadata
        """
        file_path = self.metadata_dir / metadata_file
        
        if not file_path.exists():
            raise FileNotFoundError(f"Metadata file not found: {file_path}")
        
        with open(file_path, 'r') as f:
            return json.load(f)
    
    def get_stock_info(self, symbol: str) -> Dict:
        """
        Get information about a specific stock.
        
        Args:
            symbol: Stock symbol
            
        Returns:
            Dictionary with stock information
        """
        metadata = self.load_metadata('stock_info.json')
        
        for stock in metadata.get('stocks', []):
            if stock['symbol'] == symbol:
                return stock
        
        return {}
    
    def get_date_range(self, symbol: str) -> tuple:
        """
        Get the date range for a stock dataset.
        
        Args:
            symbol: Stock symbol
            
        Returns:
            Tuple of (start_date, end_date)
        """
        df = self.load_stock(symbol)
        return (df.index.min(), df.index.max())
    
    def merge_stocks(self, symbols: List[str], column: str = 'Close') -> pd.DataFrame:
        """
        Merge specified column from multiple stocks into a single DataFrame.
        
        Args:
            symbols: List of stock symbols
            column: Column to merge (default: 'Close')
            
        Returns:
            DataFrame with merged data
        """
        dfs = []
        
        for symbol in symbols:
            df = self.load_stock(symbol)[[column]]
            df.columns = [symbol]
            dfs.append(df)
        
        return pd.concat(dfs, axis=1)
    
    def filter_by_date(self, df: pd.DataFrame, start_date: str = None, 
                      end_date: str = None) -> pd.DataFrame:
        """
        Filter DataFrame by date range.
        
        Args:
            df: DataFrame with DateTimeIndex
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)
            
        Returns:
            Filtered DataFrame
        """
        if start_date:
            df = df[df.index >= start_date]
        if end_date:
            df = df[df.index <= end_date]
        
        return df
    
    def get_summary_stats(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Get summary statistics for the dataset.
        
        Args:
            df: Stock DataFrame
            
        Returns:
            DataFrame with summary statistics
        """
        stats = pd.DataFrame({
            'Mean': df.mean(),
            'Median': df.median(),
            'Std Dev': df.std(),
            'Min': df.min(),
            'Max': df.max(),
            '25%': df.quantile(0.25),
            '75%': df.quantile(0.75)
        })
        
        return stats.round(2)
    
    def calculate_returns(self, df: pd.DataFrame, column: str = 'Close', 
                         period: int = 1) -> pd.Series:
        """
        Calculate returns for a given column.
        
        Args:
            df: Stock DataFrame
            column: Column to calculate returns on
            period: Period for return calculation (default: 1 day)
            
        Returns:
            Series with returns
        """
        return df[column].pct_change(period)
    
    def clear_cache(self):
        """Clear the cached stock data."""
        self.stocks_cache.clear()


def load_stock_data(symbol: str, data_dir: str = 'data') -> pd.DataFrame:
    """
    Convenience function to quickly load stock data.
    
    Args:
        symbol: Stock symbol
        data_dir: Directory containing data files
        
    Returns:
        DataFrame with stock data
    """
    loader = StockDataLoader(data_dir=data_dir)
    return loader.load_stock(symbol)


if __name__ == "__main__":
    # Example usage
    print("=" * 70)
    print("Stock Market Time Series Dataset - Data Loader")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Initialize loader
    loader = StockDataLoader()
    
    # Load a single stock
    print("\n1. Loading AAPL stock data...")
    aapl = loader.load_stock('AAPL')
    print(f"   Loaded {len(aapl)} trading days")
    print(f"   Date range: {aapl.index.min()} to {aapl.index.max()}")
    print(f"\n   First few rows:")
    print(aapl.head())
    
    # Load multiple stocks
    print("\n\n2. Loading multiple stocks...")
    stocks = loader.load_multiple_stocks(['AAPL', 'GOOGL', 'MSFT'])
    print(f"   Loaded {len(stocks)} stocks")
    
    # Merge closing prices
    print("\n\n3. Merging closing prices...")
    merged = loader.merge_stocks(['AAPL', 'GOOGL', 'MSFT'], column='Close')
    print(merged.head())
    
    # Get summary statistics
    print("\n\n4. Summary statistics for AAPL:")
    stats = loader.get_summary_stats(aapl)
    print(stats)
    
    # Calculate returns
    print("\n\n5. Daily returns for AAPL:")
    returns = loader.calculate_returns(aapl)
    print(f"   Mean return: {returns.mean():.4f}")
    print(f"   Std dev: {returns.std():.4f}")
    print(f"   Sharpe ratio (annualized): {(returns.mean() / returns.std()) * (252 ** 0.5):.4f}")
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets!")
    print("=" * 70)