#region imports
from AlgorithmImports import *
#endregion

# Sample research environment for wheel strategy.

from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import decimal as d

class WheelResearchAlgorithm(QCAlgorithm):

    def Initialize(self):
        self._no_K = 10       # no of strikes around ATM => for uni selection
        self.MIN_EXPIRY = 0 # min num of days to expiration => for uni selection
        self.MAX_EXPIRY = 8 # max num of days to expiration => for uni selection
        self.MAX_DELTA = 0.3
        self.MIN_PREMIUM = 0.2
        self.ticker = "AAPL"
        self.benchmarkTicker = "SPY"
        self.SetStartDate(2021, 1, 1)
        #self.SetEndDate(2020, 8, 10)
        self.SetCash(10000)
        
        self.resolution = Resolution.Daily
        self.call, self.put, self.takeProfitTicket = None, None, None
        
        equity = self.AddEquity(self.ticker, self.resolution)
        option = self.AddOption(self.ticker, self.resolution)
        self.symbol = option.Symbol

        # set strike/expiry filter for this option chain
        #option.SetFilter(-3, +3, timedelta(30), timedelta(60))
        
        # set our strike/expiry filter for this option chain
        option.SetFilter(self.UniverseFunc) 

        # for greeks and pricer (needs some warmup) - https://github.com/QuantConnect/Lean/blob/21cd972e99f70f007ce689bdaeeafe3cb4ea9c77/Common/Securities/Option/OptionPriceModels.cs#L81
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
        
        # this is needed for Greeks calcs
        self.SetWarmUp(TimeSpan.FromDays(60))    # timedelta(7)

        # use the underlying equity as the benchmark
        # self.SetBenchmark(self.benchmarkTicker)
        self.SetBenchmark(self.benchmarkTicker)
        
    def OnData(self,slice):
        if (self.IsWarmingUp): return

        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        
        if len(option_invested) == 1: return
            
        # If we already have underlying - check if we need to sell covered call
        if self.Portfolio[self.ticker].Invested:
            self.TradeCallOption(slice) 
        else:
            self.TradePutOption(slice) 
 
    def TradePutOption(self,slice):
        for i in slice.OptionChains:
            if i.Key != self.symbol: continue
        
            chain = i.Value
            
            # filter the put options contracts
            puts = [x for x in chain if x.Right == OptionRight.Put and abs(x.Greeks.Delta) > 0 and abs(x.Greeks.Delta) < self.MAX_DELTA and x.BidPrice > self.MIN_PREMIUM] 
            
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(puts, key = lambda x: x.BidPrice, reverse=True), 
                                            key = lambda x: x.Expiry)
                                            
            if len(contracts) == 0: continue  
        
            self.put = contracts[0].Symbol
            
            # short the put options
            ticket = self.MarketOrder(self.put, -1, asynchronous = False)
    
    def TradeCallOption(self,slice):
        for i in slice.OptionChains:
            if i.Key != self.symbol: continue
        
            chain = i.Value
            
            # filter the put options contracts
            calls = [x for x in chain if x.Right == OptionRight.Call and abs(x.Greeks.Delta) > 0 and abs(x.Greeks.Delta) < self.MAX_DELTA and x.BidPrice > self.MIN_PREMIUM] 
            
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(calls, key = lambda x: x.BidPrice, reverse=True), 
                                             key = lambda x: x.Expiry)

            if len(contracts) == 0: continue  
        
            self.call = contracts[0].Symbol
            
            # short the call options
            ticket = self.MarketOrder(self.call, -1, asynchronous = False)
    
 
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))
        
    def OnAssignmentOrderEvent(self, assignmentEvent):
        if self.takeProfitTicket != None:
            self.takeProfitTicket.cancel();
            self.takeProfitTicket = None

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                        .Strikes(-self._no_K, self._no_K)\
                        .Expiration(timedelta(self.MIN_EXPIRY), timedelta(self.MAX_EXPIRY))
                        
    def OnFrameworkData(self):
        return