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):

        spreadCost        = theShortContract.BidPrice - theLongContract.AskPrice
        creditReceived    = spreadCost * 100
        
        spreadWidth       = theShortContract.Strike - theLongContract.Strike
        marginReqPerSpead = spreadWidth * 1 * 100 

        requiredOrderfees = 2
        buyingPowerEffect = creditReceived - (marginReqPerSpead - requiredOrderfees)  

        marginRemaining   = self.Portfolio.MarginRemaining
        qtyWeCanPurchase  = math.floor(abs(marginRemaining / buyingPowerEffect))


        self.Debug(f"spreadWidth       = {round(theShortContract.Strike,2)} - {round(theLongContract.Strike,2)} = {round(spreadWidth,2)} ")
        self.Debug(f"marginReqPerSpead = {round(spreadWidth,2)} * 1 * 100") 
        self.Debug(f"----------")
        self.Debug(f"spreadCost        = Short BidPrice - Long AskPrice")
        self.Debug(f"spreadCost        = \t {round(theShortContract.BidPrice,2)} - {round(theLongContract.AskPrice,2)} = {round(spreadCost,2)} ")
        self.Debug(f"creditReceived    = {round(spreadCost,2)} * 100 = {round(creditReceived,2)} ")
        self.Debug(f"----------")
        self.Debug(f"requiredOrderfees = 2")
        self.Debug(f"buyingPowerEffect = {round(creditReceived,2)} - ({round(marginReqPerSpead,2)} - {requiredOrderfees}) = {round(buyingPowerEffect,2)}")  
        self.Debug(f"----------")
        self.Debug(f"marginRemaining   = {round(self.Portfolio.MarginRemaining,2)}")
        self.Debug(f"qtyWeCanPurchase  = math.floor(abs({round(marginRemaining,2)} / {round(buyingPowerEffect,2)})) = {qtyWeCanPurchase}")

        return qtyWeCanPurchase