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]

            qtyToPurchase = self.GetAffordableQty(theShortContract, theLongContract)
            
            strategy = OptionStrategies.BullPutSpread(self.optionSymbol, theShortContract.Strike, theLongContract.Strike, latest_expiry)
            
            self.Order(strategy, qtyToPurchase)
            
            self.Quit()

            
    # ===========================================================================
    def GetAffordableQty(self, theShortContract, theLongContract):
        quantities = []
        
        marginRemaining   = self.Portfolio.MarginRemaining
        for contract in [theShortContract, theLongContract]:
            security = self.Securities[contract.Symbol]
            marginParameters = InitialMarginParameters(security, 1) 
            marginRequirement = security.BuyingPowerModel.GetInitialMarginRequirement(marginParameters).Value
            affordable_quantity = int(marginRemaining / abs(marginRequirement))
            quantities.append(affordable_quantity)
                
        return min(quantities)