from QuantConnect.Securities.Option import OptionPriceModels


class SleepyFluorescentOrangeMonkey(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 1)
        self.SetEndDate(2021, 2, 5)
        
        equity = self.AddEquity("SPY", Resolution.Minute)
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
       
        option.SetFilter(self.UniverseFunc)
        #option.SetFilter(-2,2, timedelta(0), timedelta(180))
        
        self.SetWarmUp(30, Resolution.Daily)
        self.day = 0
        
        self.exp = 0
        self.strike = 0
        
        self.lors = None
        self.calls = None
        
    def UniverseFunc(self, universe):
       # include weekly contracts
       return universe.IncludeWeeklys().Expiration(TimeSpan.FromDays(2),
                                                   TimeSpan.FromDays(40)).Strikes(-30,30)

    def OnData(self, data):
        #if not data.ContainsKey(self.vx1):
        #    return
        if self.IsWarmingUp: return
        if self.day == self.Time.day:
            return
        self.day = self.Time.day
        
        if(data.HasData) :

            for kvp in data.OptionChains.Values:
                chains = kvp
                
                if kvp.Contracts.Count < 1: continue
            
                calls = list(filter(lambda x: x.Right == OptionRight.Call, chains))
                puts = list(filter(lambda x: x.Right == OptionRight.Put, chains))
                #self.underlying_price = self.Securities[self.spy].Price
                self.underlying_price = self.Securities[self.symbol].Price
                expiries = [i.Expiry for i in puts]
                strikes = [i.Strike for i in puts]

                if len(expiries) < 2:
                    return
              
                strike = min(strikes, key=lambda x: abs(x-self.underlying_price * float(0.97) ))
                self.atm_call = [i for i in chains if i.Expiry == expiries[0] and i.Strike == strike]
                self.atm_put = [i for i in chains if i.Expiry ==expiries[0] and i.Strike ==strike]
                
                calls = [i for i in chains if i.Expiry == expiries[1] and i.Strike == strike]
                puts = [i for i in chains if i.Expiry ==expiries[1] and i.Strike ==strike]
        
                if len(self.atm_call) <1 or len(calls) < 1:
                    return
                if calls[0].ImpliedVolatility == 0.0:
                    return

                self.first = self.atm_call[0].ImpliedVolatility
                self.second = calls[0].ImpliedVolatility
                
                self.Log(f"First: {self.atm_call[0].Symbol} ({self.first});  Second: {calls[0].Symbol} ({self.second})")