import pandas as pd

class ThreeATREMAIndicator:
    def __init__(self, ema_length: int = 21):
        """
        Initialize Three ATR EMA Indicator
        
        Args:
            ema_length (int): Length for EMA calculation (default: 21)
        """
        self.ema_length = ema_length

    def calculate_atr(self, high: pd.Series, low: pd.Series, close: pd.Series) -> pd.Series:
        """Calculate Average True Range"""
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        return tr

    def calculate(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Calculate Three ATR EMA indicator
        
        Args:
            df (pd.DataFrame): DataFrame with 'high', 'low', 'close' columns
            
        Returns:
            pd.DataFrame: DataFrame with added columns for EMA and ATR bands
        """
        # Calculate EMA
        ema = df['close'].ewm(span=self.ema_length, adjust=False).mean()
        
        # Calculate ATR and bands
        atr = self.calculate_atr(df['high'], df['low'], df['close'])
        upper_band = ema + (atr * 3)
        lower_band = ema - (atr * 3)
        
        # Fix signal condition with proper pandas syntax
        bullish_condition = (
            (df['close'] < ema) & 
            (df['close'].shift(1) <= lower_band.shift(1)) & 
            (df['close'] > df['close'].shift(1))
        )
        
        return pd.DataFrame({
            'ema': ema,
            'upper_band': upper_band,
            'lower_band': lower_band,
            'bullish_signal': bullish_condition  # Match the column name used in screener
        })

    def get_signals(self, df: pd.DataFrame) -> dict:
        """
        Get the current trading signals based on Three ATR EMA
        
        Args:
            df (pd.DataFrame): DataFrame with price data
            
        Returns:
            dict: Dictionary containing current signal and band values
        """
        results = self.calculate(df)
        
        # Get latest values
        if not results.empty:
            current = results.iloc[-1]
            return {
                'ema': current['ema'],
                'upper_band': current['upper_band'],
                'lower_band': current['lower_band'],
                'is_signal': current['bullish_signal']
            }
        else:
            return {}