from AlgorithmImports import *

class IndexOptionsDataAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 2, 1)
        self.SetCash(200000)

        index = self.AddIndex("SPX")
        option = self.AddOption(index.Symbol)
        option.SetFilter(-2, +2, 0, 180)
        self.optionSymbol = option.Symbol

    def OnData(self, data):
        if not self.Portfolio.Invested and self.IsMarketOpen(self.optionSymbol):
            chain = data.OptionChains.get(self.optionSymbol)
            if not chain:
                return
            
            callContracts = [c for c in chain if c.Right == OptionRight.Call]
            if callContracts:
                    
                expiry = max([c.Expiry for c in callContracts])
                callContracts = sorted([c for c in callContracts if c.Expiry == expiry],
                    key=lambda c: c.Strike)
                        
                if len(callContracts) < 2:
                    return
                
                longCall, shortCall = callContracts[0:2]
                
                # Use all the buying power
                quantity = min([
                    abs(self.CalculateOrderQuantity(shortCall.Symbol, -1)),
                    abs(self.CalculateOrderQuantity(longCall.Symbol, 1))])
                
                self.MarketOrder(shortCall.Symbol, -quantity);
                self.MarketOrder(longCall.Symbol, quantity);
                
                expectedMarginUsage = max((longCall.Strike - shortCall.Strike) * self.Securities[longCall.Symbol].SymbolProperties.ContractMultiplier * quantity, 0)
                if expectedMarginUsage != self.Portfolio.TotalMarginUsed:
                    raise Exception("Unexpect margin used!")


    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Type == SecurityType.IndexOption:
                # Historical data
                history = self.History(security.Symbol, 10, Resolution.Daily)
                self.Debug(f"We got {len(history)} from our history request for {security.Symbol}")