"""
Stock Market Time Series Dataset - Advanced Technical Indicators

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
from typing import Tuple
import warnings
warnings.filterwarnings('ignore')


class AdvancedIndicators:
    """
    Advanced technical indicators beyond basic MA, RSI, MACD.
    """
    
    @staticmethod
    def bollinger_bands(df: pd.DataFrame, window: int = 20, num_std: float = 2) -> pd.DataFrame:
        """
        Calculate Bollinger Bands.
        
        Args:
            df: DataFrame with price data
            window: Moving average window
            num_std: Number of standard deviations
            
        Returns:
            DataFrame with Bollinger Bands columns
        """
        df = df.copy()
        df['BB_Middle'] = df['Close'].rolling(window=window).mean()
        rolling_std = df['Close'].rolling(window=window).std()
        df['BB_Upper'] = df['BB_Middle'] + (rolling_std * num_std)
        df['BB_Lower'] = df['BB_Middle'] - (rolling_std * num_std)
        df['BB_Width'] = df['BB_Upper'] - df['BB_Lower']
        df['BB_PercentB'] = (df['Close'] - df['BB_Lower']) / (df['BB_Upper'] - df['BB_Lower'])
        return df
    
    @staticmethod
    def average_true_range(df: pd.DataFrame, window: int = 14) -> pd.Series:
        """
        Calculate Average True Range (ATR).
        
        Args:
            df: DataFrame with OHLC data
            window: ATR period
            
        Returns:
            Series with ATR values
        """
        high_low = df['High'] - df['Low']
        high_close = np.abs(df['High'] - df['Close'].shift())
        low_close = np.abs(df['Low'] - df['Close'].shift())
        
        ranges = pd.concat([high_low, high_close, low_close], axis=1)
        true_range = np.max(ranges, axis=1)
        atr = true_range.rolling(window).mean()
        
        return atr
    
    @staticmethod
    def stochastic_oscillator(df: pd.DataFrame, k_window: int = 14, d_window: int = 3) -> Tuple[pd.Series, pd.Series]:
        """
        Calculate Stochastic Oscillator.
        
        Args:
            df: DataFrame with OHLC data
            k_window: %K period
            d_window: %D period (SMA of %K)
            
        Returns:
            Tuple of (%K, %D) Series
        """
        low_min = df['Low'].rolling(window=k_window).min()
        high_max = df['High'].rolling(window=k_window).max()
        
        k_percent = 100 * (df['Close'] - low_min) / (high_max - low_min)
        d_percent = k_percent.rolling(window=d_window).mean()
        
        return k_percent, d_percent
    
    @staticmethod
    def commodity_channel_index(df: pd.DataFrame, window: int = 20) -> pd.Series:
        """
        Calculate Commodity Channel Index (CCI).
        
        Args:
            df: DataFrame with OHLC data
            window: CCI period
            
        Returns:
            Series with CCI values
        """
        typical_price = (df['High'] + df['Low'] + df['Close']) / 3
        sma = typical_price.rolling(window=window).mean()
        mean_deviation = typical_price.rolling(window=window).apply(
            lambda x: np.abs(x - x.mean()).mean()
        )
        cci = (typical_price - sma) / (0.015 * mean_deviation)
        
        return cci
    
    @staticmethod
    def on_balance_volume(df: pd.DataFrame) -> pd.Series:
        """
        Calculate On-Balance Volume (OBV).
        
        Args:
            df: DataFrame with Close and Volume
            
        Returns:
            Series with OBV values
        """
        obv = (np.sign(df['Close'].diff()) * df['Volume']).fillna(0).cumsum()
        return obv
    
    @staticmethod
    def money_flow_index(df: pd.DataFrame, window: int = 14) -> pd.Series:
        """
        Calculate Money Flow Index (MFI).
        
        Args:
            df: DataFrame with OHLCV data
            window: MFI period
            
        Returns:
            Series with MFI values
        """
        typical_price = (df['High'] + df['Low'] + df['Close']) / 3
        money_flow = typical_price * df['Volume']
        
        positive_flow = pd.Series(0.0, index=df.index)
        negative_flow = pd.Series(0.0, index=df.index)
        
        positive_flow[typical_price > typical_price.shift(1)] = money_flow[typical_price > typical_price.shift(1)]
        negative_flow[typical_price < typical_price.shift(1)] = money_flow[typical_price < typical_price.shift(1)]
        
        positive_mf = positive_flow.rolling(window=window).sum()
        negative_mf = negative_flow.rolling(window=window).sum()
        
        mfi = 100 - (100 / (1 + positive_mf / negative_mf))
        
        return mfi
    
    @staticmethod
    def williams_percent_r(df: pd.DataFrame, window: int = 14) -> pd.Series:
        """
        Calculate Williams %R.
        
        Args:
            df: DataFrame with OHLC data
            window: Period
            
        Returns:
            Series with Williams %R values
        """
        high_max = df['High'].rolling(window=window).max()
        low_min = df['Low'].rolling(window=window).min()
        
        wr = -100 * (high_max - df['Close']) / (high_max - low_min)
        
        return wr
    
    @staticmethod
    def parabolic_sar(df: pd.DataFrame, af_start: float = 0.02, 
                     af_increment: float = 0.02, af_max: float = 0.2) -> pd.Series:
        """
        Calculate Parabolic SAR (simplified version).
        
        Args:
            df: DataFrame with OHLC data
            af_start: Initial acceleration factor
            af_increment: AF increment
            af_max: Maximum AF
            
        Returns:
            Series with SAR values
        """
        high = df['High']
        low = df['Low']
        close = df['Close']
        
        sar = close.copy()
        sar.iloc[0] = close.iloc[0]
        
        is_uptrend = True
        af = af_start
        ep = high.iloc[0]  # Extreme point
        
        for i in range(1, len(df)):
            if is_uptrend:
                sar.iloc[i] = sar.iloc[i-1] + af * (ep - sar.iloc[i-1])
                
                if low.iloc[i] < sar.iloc[i]:
                    is_uptrend = False
                    sar.iloc[i] = ep
                    ep = low.iloc[i]
                    af = af_start
                else:
                    if high.iloc[i] > ep:
                        ep = high.iloc[i]
                        af = min(af + af_increment, af_max)
            else:
                sar.iloc[i] = sar.iloc[i-1] - af * (sar.iloc[i-1] - ep)
                
                if high.iloc[i] > sar.iloc[i]:
                    is_uptrend = True
                    sar.iloc[i] = ep
                    ep = high.iloc[i]
                    af = af_start
                else:
                    if low.iloc[i] < ep:
                        ep = low.iloc[i]
                        af = min(af + af_increment, af_max)
        
        return sar
    
    @staticmethod
    def ichimoku_cloud(df: pd.DataFrame) -> pd.DataFrame:
        """
        Calculate Ichimoku Cloud.
        
        Args:
            df: DataFrame with OHLC data
            
        Returns:
            DataFrame with Ichimoku components
        """
        df = df.copy()
        
        # Tenkan-sen (Conversion Line): (9-period high + 9-period low)/2
        period9_high = df['High'].rolling(window=9).max()
        period9_low = df['Low'].rolling(window=9).min()
        df['Ichimoku_Tenkan'] = (period9_high + period9_low) / 2
        
        # Kijun-sen (Base Line): (26-period high + 26-period low)/2
        period26_high = df['High'].rolling(window=26).max()
        period26_low = df['Low'].rolling(window=26).min()
        df['Ichimoku_Kijun'] = (period26_high + period26_low) / 2
        
        # Senkou Span A (Leading Span A): (Conversion Line + Base Line)/2
        df['Ichimoku_SpanA'] = ((df['Ichimoku_Tenkan'] + df['Ichimoku_Kijun']) / 2).shift(26)
        
        # Senkou Span B (Leading Span B): (52-period high + 52-period low)/2
        period52_high = df['High'].rolling(window=52).max()
        period52_low = df['Low'].rolling(window=52).min()
        df['Ichimoku_SpanB'] = ((period52_high + period52_low) / 2).shift(26)
        
        # Chikou Span (Lagging Span): Current closing price shifted back 26 periods
        df['Ichimoku_Chikou'] = df['Close'].shift(-26)
        
        return df
    
    @staticmethod
    def add_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
        """
        Add all advanced indicators to DataFrame.
        
        Args:
            df: DataFrame with OHLCV data
            
        Returns:
            DataFrame with all indicators
        """
        df = df.copy()
        
        # Bollinger Bands
        df = AdvancedIndicators.bollinger_bands(df)
        
        # ATR
        df['ATR'] = AdvancedIndicators.average_true_range(df)
        
        # Stochastic
        df['Stoch_K'], df['Stoch_D'] = AdvancedIndicators.stochastic_oscillator(df)
        
        # CCI
        df['CCI'] = AdvancedIndicators.commodity_channel_index(df)
        
        # OBV
        df['OBV'] = AdvancedIndicators.on_balance_volume(df)
        
        # MFI
        df['MFI'] = AdvancedIndicators.money_flow_index(df)
        
        # Williams %R
        df['Williams_R'] = AdvancedIndicators.williams_percent_r(df)
        
        # Parabolic SAR
        df['SAR'] = AdvancedIndicators.parabolic_sar(df)
        
        # Ichimoku Cloud
        df = AdvancedIndicators.ichimoku_cloud(df)
        
        return df


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Advanced Indicators")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Load data
    df = pd.read_csv('data/AAPL.csv', parse_dates=['Date'], index_col='Date')
    
    print(f"\nOriginal columns: {list(df.columns)}")
    print(f"Number of columns: {len(df.columns)}")
    
    # Add all indicators
    df_enhanced = AdvancedIndicators.add_all_indicators(df)
    
    print(f"\nEnhanced columns: {list(df_enhanced.columns)}")
    print(f"Number of columns: {len(df_enhanced.columns)}")
    print(f"\nNew indicators added: {len(df_enhanced.columns) - len(df.columns)}")
    
    # Show sample data
    print(f"\nSample data with new indicators:")
    print(df_enhanced[['Close', 'BB_Upper', 'BB_Lower', 'ATR', 'Stoch_K', 
                       'CCI', 'MFI', 'Williams_R']].tail())
    
    # Save enhanced dataset
    output_path = 'data/AAPL_enhanced.csv'
    df_enhanced.to_csv(output_path)
    print(f"\nEnhanced dataset saved to: {output_path}")
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)