from QuantConnect.Securities.Option import OptionPriceModels

class BasicTemplateOptionsConsolidationAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2013, 10, 7)
        self.SetEndDate(2013, 10, 8)
        self.SetCash(1000000)

        self.consolidation_period = 60

        self.option = self.AddOption('SPY')
        self.option.SetFilter(-1, +1, 0, 30)
        self.option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        self.SetWarmUp(30, Resolution.Daily)
        
        self.consolidators = dict()
    
    def OnData(self, data):
        for symbol, chain in data.OptionChains.items():
            for contract in chain:
                self.consolidators[contract.Symbol].Update(IndicatorDataPoint(self.Time, contract.ImpliedVolatility))
            

    def OnIVConsolidated(self, sender, bar):
        self.Log(f"Consolidated IV for {bar.Symbol} received at {self.Time} ... O:{bar.Open}; H:{bar.High}; L:{bar.Low}; C:{bar.Close}")
                    
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Type == SecurityType.Option:
                symbol = security.Symbol
                consolidator = BaseDataConsolidator(timedelta(hours=1))
                consolidator.DataConsolidated += self.OnIVConsolidated
                self.consolidators[symbol] = consolidator
                
        for security in changes.RemovedSecurities:
            if security.Type == SecurityType.Option:
                consolidator = self.consolidators.pop(security.Symbol)
                consolidator.DataConsolidated -= self.OnIVConsolidated