from AlgorithmImports import *
from datetime import datetime, timedelta

class MonthlyGapTradingAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 7, 3)
        self.SetEndDate(2024, 7, 1)
        self.SetCash(10000000)
        
        #self.UniverseSettings.Resolution = Resolution.Daily

        # Use SPY constituents as the universe
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.AddUniverseSelection(ETFConstituentsUniverseSelectionModel(self.spy))

        # self.aapl = self.AddEquity("AAPL", Resolution.Daily).Symbol
        # self.msft = self.AddEquity("MSFT", Resolution.Daily).Symbol
        # self.googl = self.AddEquity("GOOGL", Resolution.Daily).Symbol
        # self.AddUniverse(self.Universe.ETF(self.aapl))
        # self.AddUniverse(self.Universe.ETF(self.msft))


        self.atr_period = 14
        
        self.data = {}
        self.atr = {}
        self.stop_levels = {}

        self.buy_threshold = 9
        self.sell_threshold = 9

        self.long_positions = set()
        self.short_positions = set()

        #self.SetBenchmark("SPY")

        #self.SetSecurityInitializer(lambda x: x.SetFeeModel(ConstantFeeModel(0)))

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol.SecurityType == SecurityType.Equity and self.Securities.ContainsKey(symbol):
                try:
                    self.InitializeSymbolData(symbol)
                except Exception as e:
                    self.Log(f"Error initializing {symbol}: {str(e)}")

        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.data:
                if self.data[symbol]['consolidator'] is not None:
                    self.SubscriptionManager.RemoveConsolidator(symbol, self.data[symbol]['consolidator'])
                del self.data[symbol]
            if symbol in self.atr:
                del self.atr[symbol]

    def InitializeSymbolData(self, symbol):
        if symbol not in self.data:
            self.data[symbol] = {
                'buy_count': 0,
                'sell_count': 0,
                'consolidator': None,
                'price_window': RollingWindow[float](5)
            }
            
            # Only create and register the consolidator if the symbol is in Securities
            if self.Securities.ContainsKey(symbol):
                self.data[symbol]['consolidator'] = self.Consolidate(symbol, CalendarType.Monthly, self.OnMonthlyBar)
            else:
                self.Log(f"Warning: Symbol {symbol} not found in Securities. Skipping consolidator creation.")
        
        if symbol not in self.atr and self.Securities.ContainsKey(symbol):
            self.atr[symbol] = self.ATR(symbol, self.atr_period, Resolution.Daily)

    def OnMonthlyBar(self, bar):
        symbol = bar.Symbol
        
        # Check if the symbol data exists, if not, initialize it
        if symbol not in self.data:
            self.InitializeSymbolData(symbol)
        
        symbol_data = self.data[symbol]

        symbol_data['price_window'].Add(bar.Close)

        if symbol_data['price_window'].IsReady:
            self.EvaluatePosition(symbol)

        self.UpdateStopLosses(symbol)

    def EvaluatePosition(self, symbol):
        symbol_data = self.data[symbol]
        current_close = symbol_data['price_window'][0]
        five_months_ago_close = symbol_data['price_window'][4]

        if current_close > five_months_ago_close:
            symbol_data['buy_count'] += 1
            symbol_data['sell_count'] = 0
        elif current_close < five_months_ago_close:
            symbol_data['sell_count'] += 1
            symbol_data['buy_count'] = 0
        else:
            symbol_data['buy_count'] = 0
            symbol_data['sell_count'] = 0

        if symbol_data['buy_count'] > self.buy_threshold and symbol not in self.short_positions:
            self.short_positions.add(symbol)
            #self.Order(symbol, 1000 , tag="signal_tag") 

            if symbol in self.long_positions:
                self.long_positions.remove(symbol)

        if symbol_data['sell_count'] > self.sell_threshold and symbol not in self.long_positions:
            self.long_positions.add(symbol)
            #self.Order(symbol, -1000 , tag="signal_tag")

            if symbol in self.short_positions:
                self.short_positions.remove(symbol)

        self.Rebalance()

    # def Rebalance(self):
    #     num_longs = len(self.long_positions)
    #     num_shorts = len(self.short_positions)
    #     max_allocation = 0.05

    #     if num_longs > 0:
    #         long_allocation = min(max_allocation, 1.0 / num_longs)
    #         for symbol in self.long_positions:
    #             self.SetHoldings(symbol, long_allocation)

    #     if num_shorts > 0:
    #         short_allocation = min(max_allocation, 1.0 / num_shorts)
    #         for symbol in self.short_positions:
    #             self.SetHoldings(symbol, -short_allocation)
    

    def Rebalance(self):
        num_longs = len(self.long_positions)
        num_shorts = len(self.short_positions)
        max_allocation = 0.05

        if num_longs > 0:
            long_allocation = min(max_allocation, 1.0 / num_longs)
            for symbol in self.long_positions:
                current_holding = self.Portfolio[symbol].Quantity
                security_price = self.Securities[symbol].Price
                
                if security_price > 0:
                    if current_holding == 0:
                        self.SetHoldings(symbol, long_allocation, tag="New Position Entered")
                    else:
                        self.SetHoldings(symbol, long_allocation , tag = "Rebalance Position" )
                else:
                    self.Log(f"Warning: Zero or negative price for {symbol}. Skipping rebalance for this symbol.")

        if num_shorts > 0:
            short_allocation = min(max_allocation, 1.0 / num_shorts)
            for symbol in self.short_positions:
                current_holding = self.Portfolio[symbol].Quantity
                security_price = self.Securities[symbol].Price
                
                if security_price > 0:
                    if current_holding == 0:
                        self.SetHoldings(symbol, -short_allocation, tag="New Position Entered")
                    else:
                        self.SetHoldings(symbol, -short_allocation , tag = "Rebalance Position")
                else:
                    self.Log(f"Warning: Zero or negative price for {symbol}. Skipping rebalance for this symbol.")

    def UpdateStopLosses(self, symbol):
        if symbol not in self.data or symbol not in self.atr or not self.Securities.ContainsKey(symbol):
            return

        current_price = self.Securities[symbol].Price
        if current_price <= 0:
            self.Log(f"Warning: Zero or negative price for {symbol}. Skipping stop loss update for this symbol.")
            return

        atr_value = self.atr[symbol].Current.Value

        if symbol in self.long_positions:
            stop_price = current_price - 2 * atr_value
            if symbol not in self.stop_levels or stop_price > self.stop_levels[symbol]:
                self.stop_levels[symbol] = stop_price
                self.StopMarketOrder(symbol, -self.Portfolio[symbol].Quantity, stop_price , tag = "Stop Order")

        elif symbol in self.short_positions:
            stop_price = current_price + 2 * atr_value
            if symbol not in self.stop_levels or stop_price < self.stop_levels[symbol]:
                self.stop_levels[symbol] = stop_price
                self.StopMarketOrder(symbol, -self.Portfolio[symbol].Quantity, stop_price , tag = "Stop Order")

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            if order.Type == OrderType.StopMarket:
                symbol = order.Symbol
                if symbol in self.long_positions:
                    self.long_positions.remove(symbol)
                elif symbol in self.short_positions:
                    self.short_positions.remove(symbol)

        self.Log(f"Date: {self.Time}, Order Event: {orderEvent}")

    def OnEndOfMonth(self):
        self.Log(f"Date: {self.Time}, Long Positions: {self.long_positions}, Short Positions: {self.short_positions}")