import pandas as pd

class BollingerBands:
    """
    Calculates Bollinger Bands for given financial data.
    """
    def __init__(self, config):
        """
        Initializes the BollingerBands calculator.

        Args:
            period (int): The period for the moving average and standard deviation.
            std_dev_multiplier (float): The number of standard deviations for the upper and lower bands.
            bb_width (float): The width of the Bollinger Bands.
        """
        if config['bb_period'] <= 0:
            raise ValueError("Period must be a positive integer.")
        if config['trending']['bb_std_dev_multiplier'] <= 0 or config['sideways']['bb_std_dev_multiplier'] <= 0:
            raise ValueError("Standard deviation multiplier must be positive.")
        if config['bb_width'] <= 0:
            raise ValueError("BB width must be positive.")
            
        self.config = config

    def calculate(self, data_df: pd.DataFrame, price_column: str = 'close', squeeze = False) -> pd.DataFrame:
        """
        Calculates Bollinger Bands and adds them to the DataFrame.

        Args:
            data_df (pd.DataFrame): DataFrame with price data. Must include the price_column.
            price_column (str): The name of the column containing the price data (e.g., 'close').

        Returns:
            pd.DataFrame: The original DataFrame with added columns:
                          'SMA' (Simple Moving Average),
                          'UpperBand',
                          'LowerBand'.
        """
        if price_column not in data_df.columns:
            raise ValueError(f"Price column '{price_column}' not found in DataFrame.")

        if not squeeze:
            # Calculate SMA
            data_df['SMA'] = data_df[price_column].rolling(window=self.config['bb_period']).mean()

            # Calculate Standard Deviation
            std_dev = data_df[price_column].rolling(window=self.config['bb_period']).std()

            # Calculate Upper and Lower Bands
            data_df['UpperBand'] = data_df['SMA'] + (2.0* std_dev)
            data_df['LowerBand'] = data_df['SMA'] - (2.0* std_dev)

            # Calculate the width of the Bollinger Bands
            data_df['BBWidth'] = (data_df['UpperBand'] - data_df['LowerBand']) / data_df['SMA']

            # Calculate the market regime
            # 1 = sideways, 0 = trending
            data_df['MarketRegime'] = (data_df['BBWidth'] < self.config['bb_width']).astype(int)

            if data_df['MarketRegime'].sum() > 0:
                data_df['UpperBand'] = data_df['SMA'] + (self.config['trending']['bb_std_dev_multiplier'] * std_dev)
                data_df['LowerBand'] = data_df['SMA'] - (self.config['trending']['bb_std_dev_multiplier'] * std_dev)
            else:
                data_df['UpperBand'] = data_df['SMA'] + (self.config['sideways']['bb_std_dev_multiplier'] * std_dev)
                data_df['LowerBand'] = data_df['SMA'] - (self.config['sideways']['bb_std_dev_multiplier'] * std_dev)

        else:
            data_df['SMA'] = data_df[price_column].rolling(window=14).mean()
            # Calculate Standard Deviation
            std_dev = data_df[price_column].rolling(window=14).std()
            # Calculate Upper and Lower Bands
            data_df['UpperBand'] = data_df['SMA'] + 1.5* std_dev
            data_df['LowerBand'] = data_df['SMA'] - 1.5* std_dev

        return data_df