class QuantumParticleAtmosphericScrubbers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 12, 8)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.spy_symbol = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.Securities[self.spy_symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        self.straddleTickets = {}

    def OnData(self, data):
        
        if not data.Bars.ContainsKey(self.spy_symbol):
            return
        
        
        if not self.Portfolio.Invested:
            
            price = data.Bars[self.spy_symbol].Close;
            expiry = self.Time + timedelta(days = 14)
            
            call_symbol, put_symbol = self.GetStraddleContracts(self.spy_symbol, price, expiry)
            
            call = self.AddOptionContract(call_symbol, Resolution.Minute)
            put = self.AddOptionContract(put_symbol, Resolution.Minute)
            
            call_tickets, put_tickets = self.TradeStraddle(call, put, 2)
        
            self.straddleTickets[self.spy_symbol] = [call_tickets, put_tickets]
            
            
    def GetStraddleContracts(self, symbol, strike, expiry):
        
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
        
        expiry_sorted = sorted(contracts, key=lambda k: abs(k.ID.Date - expiry), reverse=False)
        closest_expiry = expiry_sorted[0].ID.Date
        
        contracts_with_desired_expiry = [symbol for symbol in contracts if symbol.ID.Date == closest_expiry]
        
        calls = [symbol for symbol in contracts_with_desired_expiry if symbol.ID.OptionRight == OptionRight.Call]
        
        puts = [symbol for symbol in contracts_with_desired_expiry if symbol not in calls]
        
        sorted_calls = sorted(calls, key=lambda k: abs(k.ID.StrikePrice - strike), reverse=False)
        
        sorted_puts = sorted(puts, key=lambda k: abs(k.ID.StrikePrice - strike), reverse=False)
        
        return sorted_calls[0], sorted_puts[0]
        
        
    def TradeStraddle(self, call, put, quantity):
        
        call_entry= call.AskPrice
        call_entry_ticket = self.MarketOrder(call.Symbol, quantity)
        
        call_take_profit = call_entry * 1.10
        call_stop_loss = call_entry * 0.90
        
        call_profit_ticket = self.LimitOrder(call.Symbol, -quantity, call_take_profit)
        call_loss_ticket = self.StopMarketOrder(call.Symbol, -quantity, call_stop_loss)
        
        put_entry = put.AskPrice
        put_entry_ticket = self.MarketOrder(put.Symbol, quantity)
        
        put_take_profit = put_entry * 1.10
        put_stop_loss = put_entry * 0.90
        
        put_profit_ticket = self.LimitOrder(put.Symbol, -quantity, put_take_profit)
        put_loss_ticket = self.StopMarketOrder(put.Symbol, -quantity, put_stop_loss)
        
        return (call_profit_ticket, call_loss_ticket), (put_profit_ticket, put_loss_ticket)
        
    
    def CustomSecurityInitializer(self, security):
        bar = self.GetLastKnownPrice(security)
        security.SetMarketPrice(bar)
        
    def OnOrderEvent(self, orderevent):
        
        if orderevent.Status != OrderStatus.Filled:
            return
        
        orderId = orderevent.OrderId
        
        for underlying, order_pairs in self.straddleTickets.items():
            for order_pair in order_pairs:
                order_one = order_pair[0]
                order_two = order_pair[1]
                if orderId == order_one.OrderId:
                    order_two.Cancel()
                elif orderId == order_two.OrderId:
                    order_one.Cancel()