import logging
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

from cycles.utils.storage import Storage
from cycles.utils.data_utils import aggregate_to_daily
from cycles.Analysis.boillinger_band import BollingerBands
from cycles.Analysis.rsi import RSI
from cycles.Analysis.strategies import Strategy

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("backtest.log"),
        logging.StreamHandler()
    ]
)

config_minute = {
    "start_date": "2023-01-01",
    "stop_date": "2024-01-01",
    "data_file": "btcusd_1-min_data.csv"
}

config_day = {
    "start_date": "2023-01-01",
    "stop_date": "2024-01-01",
    "data_file": "btcusd_1-day_data.csv"
}

config_strategy = {
    "bb_width": 0.05,
    "bb_period": 20,
    "rsi_period": 14,
    "trending": {
        "rsi_threshold": [30, 70],
        "bb_std_dev_multiplier": 2.5,
    },
    "sideways": {
        "rsi_threshold": [40, 60],
        "bb_std_dev_multiplier": 1.8,
    },
    "strategy_name": "MarketRegimeStrategy",
    "SqueezeStrategy": True
}

IS_DAY = False

if __name__ == "__main__":

    storage = Storage(logging=logging)

    if IS_DAY:
        config = config_day
    else:
        config = config_minute

    data = storage.load_data(config["data_file"], config["start_date"], config["stop_date"])

    if not IS_DAY:
        data_daily = aggregate_to_daily(data)
        storage.save_data(data, "btcusd_1-day_data.csv")
        df_to_plot = data_daily
    else:
        df_to_plot = data

    bb = BollingerBands(config=config_strategy)
    data_bb = bb.calculate(df_to_plot.copy())

    rsi_calculator = RSI(config=config_strategy)
    data_with_rsi = rsi_calculator.calculate(df_to_plot.copy(), price_column='close')
    
    # Combine BB and RSI data into a single DataFrame for signal generation
    # Ensure indices are aligned; they should be as both are from df_to_plot.copy()
    if 'RSI' in data_with_rsi.columns:
        data_bb['RSI'] = data_with_rsi['RSI']
    else:
        # If RSI wasn't calculated (e.g., not enough data), create a dummy column with NaNs
        # to prevent errors later, though signals won't be generated.
        data_bb['RSI'] = pd.Series(index=data_bb.index, dtype=float)
        logging.warning("RSI column not found or not calculated. Signals relying on RSI may not be generated.")

    strategy = Strategy(config=config_strategy)
    buy_condition, sell_condition = strategy.run(data_bb, config_strategy["strategy_name"])

    buy_signals = data_bb[buy_condition]
    sell_signals = data_bb[sell_condition]
    
    # plot the data with seaborn library
    if df_to_plot is not None and not df_to_plot.empty:
        # Create a figure with two subplots, sharing the x-axis
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(16, 8), sharex=True)

        # Plot 1: Close Price and Bollinger Bands
        sns.lineplot(x=data_bb.index, y='close', data=data_bb, label='Close Price', ax=ax1)
        sns.lineplot(x=data_bb.index, y='UpperBand', data=data_bb, label='Upper Band (BB)', ax=ax1)
        sns.lineplot(x=data_bb.index, y='LowerBand', data=data_bb, label='Lower Band (BB)', ax=ax1)
        # Plot Buy/Sell signals on Price chart
        if not buy_signals.empty:
            ax1.scatter(buy_signals.index, buy_signals['close'], color='green', marker='o', s=20, label='Buy Signal', zorder=5)
        if not sell_signals.empty:
            ax1.scatter(sell_signals.index, sell_signals['close'], color='red', marker='o', s=20, label='Sell Signal', zorder=5)
        ax1.set_title('Price and Bollinger Bands with Signals')
        ax1.set_ylabel('Price')
        ax1.legend()
        ax1.grid(True)

        # Plot 2: RSI
        if 'RSI' in data_bb.columns: # Check data_bb now as it should contain RSI
            sns.lineplot(x=data_bb.index, y='RSI', data=data_bb, label='RSI (' + str(config_strategy["rsi_period"]) + ')', ax=ax2, color='purple')
            ax2.axhline(config_strategy["trending"]["rsi_threshold"][1], color='red', linestyle='--', linewidth=0.8, label='Overbought (' + str(config_strategy["trending"]["rsi_threshold"][1]) + ')')
            ax2.axhline(config_strategy['trending']['rsi_threshold'][0], color='green', linestyle='--', linewidth=0.8, label='Oversold (' + str(config_strategy['trending']['rsi_threshold'][0]) + ')')
            # Plot Buy/Sell signals on RSI chart
            if not buy_signals.empty:
                ax2.scatter(buy_signals.index, buy_signals['RSI'], color='green', marker='o', s=20, label='Buy Signal (RSI)', zorder=5)
            if not sell_signals.empty:
                ax2.scatter(sell_signals.index, sell_signals['RSI'], color='red', marker='o', s=20, label='Sell Signal (RSI)', zorder=5)
            ax2.set_title('Relative Strength Index (RSI) with Signals')
            ax2.set_ylabel('RSI Value')
            ax2.set_ylim(0, 100) # RSI is typically bounded between 0 and 100
            ax2.legend()
            ax2.grid(True)
        else:
            logging.info("RSI data not available for plotting.")

        # Plot 3: BB Width
        sns.lineplot(x=data_bb.index, y='BBWidth', data=data_bb, label='BB Width', ax=ax3)
        sns.lineplot(x=data_bb.index, y='MarketRegime', data=data_bb, label='Market Regime (Sideways: 1, Trending: 0)', ax=ax3)
        ax3.set_title('Bollinger Bands Width')
        ax3.set_ylabel('BB Width')
        ax3.legend()
        ax3.grid(True)

        plt.xlabel('Date') # Common X-axis label
        fig.tight_layout() # Adjust layout to prevent overlapping titles/labels
        plt.show()
    else:
        logging.info("No data to plot.")