from QuantConnect.Securities.Option import OptionPriceModels

class AutomatonTrade(QCAlgorithm):

    def Initialize(self):
        
        # Init
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 1, 3)

        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        self.tsla = Symbol.Create("TSLA", SecurityType.Equity, Market.USA)
                
        # Live universe with selection.
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.called = 0
        
        

    def CustomSecurityInitializer(self, security):
        '''Initialize the security with raw prices'''
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        if security.Type == SecurityType.Equity:
            security.VolatilityModel = StandardDeviationOfReturnsVolatilityModel(30)
            history = self.History(security.Symbol, 31, Resolution.Daily)
            if history.empty or 'close' not in history.columns:
                return
            
            for time, row in history.loc[security.Symbol].iterrows():
                trade_bar = TradeBar(time, security.Symbol, row.open, row.high, row.low, row.close, row.volume)    
                security.VolatilityModel.Update(security, trade_bar)
        
        elif security.Type == SecurityType.Option:
            security.PriceModel = OptionPriceModels.CrankNicolsonFD() #  BlackScholes()
            

    def CoarseSelectionFunction(self, coarse):
        return [self.tsla]
        
        
    def OnData(self, data):
        
        self.called += 1
        
        symbols = self.OptionChainProvider.GetOptionContractList(self.tsla, data.Time)[:3]
        for symbol in symbols:
            if symbol not in self.Securities.Keys:
                self.AddOptionContract(symbol, Resolution.Minute)
    
        for symbol, chain in data.OptionChains.items():
            for contract in chain:
                self.Log(f"Delta {contract.Greeks.Delta} for {contract}")
                    
        if self.called >= 4:
            self.Quit()