from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Securities.Option import OptionStrategies
import math 

class QuantumVerticalInterceptor(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(2020, 1, 5)
        self.SetEndDate(2020, 1, 7)
        self.SetCash(10000)

        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        # self.SetWarmup(30, Resolution.Daily)

        # Add the option
        self.equity = self.AddEquity("SPY")
        self.option = self.AddOption("SPY")
        self.optionSymbol = self.option.Symbol

        # Add the initial contract filter
        self.option.SetFilter(-5, +5, 5, 10)


    def OnData(self,slice):
        
        for i in slice.OptionChains:
            if i.Key != self.optionSymbol: continue
            chain = i.Value
            puts = [x for x in chain if x.Right == OptionRight.Put]
            if len(puts) < 2:
                return
            
            latest_expiry = sorted(puts, key=lambda x: x.Expiry)[-1].Expiry
            puts = [x for x in puts if x.Expiry == latest_expiry]
            if len(puts) < 2:
                return
            
            sorted_by_strike = sorted(puts, key=lambda x: x.Strike)
            
            theLongContract  = sorted_by_strike[0]
            theShortContract = sorted_by_strike[1]
            
            strategy = OptionStrategies.BullPutSpread(self.optionSymbol, theShortContract.Strike, theLongContract.Strike, latest_expiry)

            self.Order(strategy, 54) # Works in interval [1, 54]
            
            #for i in range(64): # Works in interval [1, 64]
            #    self.Order(strategy, 1)
            
            self.Quit(f"Margin used: {self.Portfolio.TotalMarginUsed}; Remaining: {self.Portfolio.MarginRemaining}")