from AlgorithmImports import *
import numpy as np
import pandas as pd

class ForexStochasticBollinger(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2007, 1, 1)  # Test period start
        self.SetEndDate(2024, 12, 31)  # Test period end
        self.SetCash(100000)
        
        # List of currency pairs
        self.pairs = ["EURUSD", "GBPUSD", "USDCAD", "USDCHF", "AUDUSD", "NZDUSD", "USDJPY", "USDSEK", "USDNOK"]
        self.symbols = [self.AddForex(pair, Resolution.Daily).Symbol for pair in self.pairs]

        # Log generated symbols for debugging
        self.Debug(f"Generated Symbols: {', '.join(str(symbol) for symbol in self.symbols)}")

        # Initialize indicators and data structures
        self.indicators = {symbol: (self.STO(symbol, 14, 1, 3), self.BB(symbol, 20, 2)) for symbol in self.symbols}
        self.daily_returns = {symbol: [] for symbol in self.symbols}
        self.current_positions = {}
        self.trades_today = 0

        # Initialize transaction log
        self.transaction_log = []

        # Set warm-up period for indicators
        self.SetWarmUp(30, Resolution.Daily)

        # Schedule trading logic to run at 0:00 AM daily
        self.Schedule.On(self.DateRules.EveryDay(),
                            self.TimeRules.At(0, 00),
                            self.ExecuteStrategy)

        # Schedule liquidation event at 23:59 every day
        self.Schedule.On(self.DateRules.EveryDay(),
                            self.TimeRules.At(23, 59),
                            self.LiquidatePositions)


    def LiquidatePositions(self):
        self.Debug("Liquidating all positions at 23:59.")
        self.liquidate()


    def ExecuteStrategy(self):
        if self.IsWarmingUp:
            return

        # Reset trades count at the start of the day
        self.trades_today = 0
        
        long_signals = []
        short_signals = []

        # Evaluate signals for each pair
        for symbol in self.symbols:
            stoch, bb = self.indicators[symbol]
            current_k = stoch.StochK.Current.Value
            current_d = stoch.StochD.Current.Value
            upper_band = bb.UpperBand.Current.Value
            lower_band = bb.LowerBand.Current.Value

            # Calculate ADX
            adx_value = self.CalculateADX(symbol)

            # Skip trade if ADX
            if adx_value < 10:
                self.Debug(f"Skipping trade for {symbol} due to low ADX {adx_value:.4f}.")
                continue

            # Check long entry condition
            if current_k < 30 and current_d < 30 and self.Securities[symbol].Price < lower_band:
                long_signals.append(symbol)

            # Check short entry condition
            elif current_k > 70 and current_d > 70 and self.Securities[symbol].Price > upper_band:
                short_signals.append(symbol)

        total_signals = len(long_signals) + len(short_signals)

        # Proceed only if there are more than 2 signals
        if total_signals > 0:
            self.Debug(f"Detected Long Signals: {', '.join(str(symbol.Value) for symbol in long_signals)}")
            self.Debug(f"Detected Short Signals: {', '.join(str(symbol.Value) for symbol in short_signals)}")

        if total_signals > 2:
            # Retrieve historical prices for detected signals
            history = self.History(long_signals + short_signals, 30, Resolution.Daily)

            if history.empty:
                self.Debug("No historical data returned.")
                return

            close_prices = history['close']
            price_data = close_prices.reset_index().pivot(index='time', columns='symbol', values='close')
            daily_returns = price_data.pct_change().dropna()
            correlation_matrix = daily_returns.corr()

            # Log the correlation matrix for debugging
            self.Debug(f"Correlation Matrix:\n{correlation_matrix}")
            
            # Find the pairs with the lowest absolute correlation
            lowest_abs_corr = float('inf')
            selected_symbols = None

            # Iterate through the correlation matrix
            for i in range(len(correlation_matrix.columns)):
                for j in range(i + 1, len(correlation_matrix.columns)):
                    current_corr = correlation_matrix.iloc[i, j]
                    abs_corr = abs(current_corr)
                    if abs_corr < lowest_abs_corr:
                        lowest_abs_corr = abs_corr
                        selected_symbols = (correlation_matrix.columns[i], correlation_matrix.columns[j])

            if selected_symbols is not None:
                self.Debug(f"Selected Symbols: {selected_symbols[0]} and {selected_symbols[1]} with lowest absolute correlation {lowest_abs_corr:.4f}")
                allocation = 0.5  # Equal weight for two trades

                # Execute trades based on selected symbols
                for symbol in selected_symbols:
                    if symbol in long_signals:
                        self.Trade(symbol, True, allocation)
                    elif symbol in short_signals:
                        self.Trade(symbol, False, allocation)

            else:
                self.Debug("No pairs found with sufficient correlation.")
                return  # If no pairs found, exit

        else:
            selected_symbols = long_signals + short_signals  # Use available signals
            allocation = 1.0 / total_signals if total_signals > 0 else 0  # 100% for one signal or equal for two

            # Execute trades based on available signals
            for symbol in selected_symbols:
                if symbol in long_signals:
                    self.Trade(symbol, True, allocation)
                elif symbol in short_signals:
                    self.Trade(symbol, False, allocation)


    def Trade(self, symbol, long_signal, allocation):
        if self.trades_today < 2:
            if long_signal:
                if symbol not in self.current_positions or self.current_positions[symbol] != 1:
                    self.Liquidate(symbol)
                    self.SetHoldings(symbol, allocation)  # Set holdings based on allocation
                    self.current_positions[symbol] = 1
                    self.Debug(f"Entered Long for {symbol} with allocation {allocation:.4f}")
                    self.LogTransaction(symbol, allocation)  # Log the transaction

            else:
                if symbol not in self.current_positions or self.current_positions[symbol] != -1:
                    self.Liquidate(symbol)
                    self.SetHoldings(symbol, -allocation)  # Set holdings based on allocation
                    self.current_positions[symbol] = -1
                    self.Debug(f"Entered Short for {symbol} with allocation {allocation:.4f}")
                    self.LogTransaction(symbol, -allocation)  # Log the transaction

            self.trades_today += 1


    def LogTransaction(self, symbol, amount):
        transaction_date = self.Time  # Get the current date and time
        self.transaction_log.append({
            "date": transaction_date,  # Store the full date and time
            "pair": symbol,
            "amount": amount
        })
        self.Debug(f"Logged Transaction: {transaction_date}, {symbol}, Amount: {amount:.4f}")


    def CalculateADX(self, symbol, period=14):
        """Calculate ADX for the given symbol."""
        # Fetch historical data for ADX calculation
        history = self.History(symbol, period + 1, Resolution.Daily)
        if history.empty:
            return 0
        
        high = history['high'].values
        low = history['low'].values
        close = history['close'].values

        # Calculate the directional movements
        up_move = np.maximum(high[1:] - high[:-1], 0)
        down_move = np.maximum(low[:-1] - low[1:], 0)

        # Calculate True Range
        tr1 = high[1:] - low[1:]
        tr2 = abs(high[1:] - close[:-1])
        tr3 = abs(low[1:] - close[:-1])
        true_range = np.maximum(np.maximum(tr1, tr2), tr3)

        # Calculate the smoothed averages
        atr = np.mean(true_range[-period:])  # Average True Range
        up_movement_avg = np.mean(up_move[-period:])  # Average Up Movement
        down_movement_avg = np.mean(down_move[-period:])  # Average Down Movement

        # Calculate the Directional Index
        if up_movement_avg + down_movement_avg == 0:
            return 0
        plus_di = 100 * (up_movement_avg / atr)
        minus_di = 100 * (down_movement_avg / atr)

        # Calculate ADX
        adx = 100 * np.mean(np.abs(plus_di - minus_di) / (plus_di + minus_di)) if (plus_di + minus_di) != 0 else 0

        return adx

# Entry point for the algorithm
model = ForexStochasticBollinger()