from datetime import timedelta
from QuantConnect.Data.Custom.CBOE import *

class BasicTemplateAlgorithm(QCAlgorithm):

    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):

        self.SetStartDate(2017,1, 1)  #Set Start Date
        self.SetEndDate(2020,12,31)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.rsi = self.RSI("SPY", 14, MovingAverageType.Simple, Resolution.Minute)
        self.contract = str()
        self.contractsAdded = set()
        self.symbol = self.equity.Symbol

        self.DaysBeforeExp = 2 # number of days before expiry to exit
        self.DTE = 25 # target days till expiration
        self.OTM = 0.01 # target percentage OTM of put

        # schedule Plotting function 30 minutes after every market open
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), \
                        self.TimeRules.AfterMarketOpen(self.symbol, 30), \
                        self.Plotting)
        self.SetWarmUp(14)

    

    def OnData(self, data):

        if self.IsWarmingUp:
            return

        if self.rsi.Current.Value < 15:
                self.Debug("RSI is less than 15")
                self.BuyCall(data)
                
        if self.contract:
            if (self.contract.ID.Date - self.Time) <= timedelta(self.DaysBeforeExp):
                self.Liquidate(self.contract)
                self.Log("Closed: too close to expiration")
                self.contract = str()

 

    def BuyCall(self, data):
        # get option data
        if self.contract == str():
            self.contract = self.OptionsFilter(data)
            return

# if not invested and option data added successfully, buy option
        elif not self.Portfolio[self.contract].Invested and data.ContainsKey(self.contract):
            self.Buy(self.contract, 10)


    def OptionsFilter(self, data):

        ''' OptionChainProvider gets a list of option contracts for an underlying symbol at requested date.
            Then you can manually filter the contract list returned by GetOptionContractList.
            The manual filtering will be limited to the information included in the Symbol
            (strike, expiration, type, style) and/or prices from a History call '''


        contracts = self.OptionChainProvider.GetOptionContractList(self.symbol, data.Time)
        self.underlyingPrice = self.Securities[self.symbol].Price
        # filter the out-of-money put options from the contract list which expire close to self.DTE num of days from now
        otm_puts = [i for i in contracts if i.ID.OptionRight == OptionRight.Call and
                                            self.underlyingPrice - i.ID.StrikePrice > self.OTM * self.underlyingPrice and
                                            self.DTE - 8 < (i.ID.Date - data.Time).days < self.DTE + 8]

        if len(otm_puts) > 0:
            # sort options by closest to self.DTE days from now and desired strike, and pick first
            contract = sorted(sorted(otm_puts, key = lambda x: abs((x.ID.Date - self.Time).days - self.DTE)),
                                                     key = lambda x: self.underlyingPrice - x.ID.StrikePrice)[0]
            if contract not in self.contractsAdded:
                self.contractsAdded.add(contract)
                # use AddOptionContract() to subscribe the data for specified contract
                self.AddOptionContract(contract, Resolution.Minute)
            return contract

        else:
            return str()


    def Plotting(self):
        # plot underlying's price
        self.Plot("Data Chart", self.symbol, self.Securities[self.symbol].Close)
        # plot strike of put option

        
        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        if option_invested:
                self.Plot("Data Chart", "strike", option_invested[0].ID.StrikePrice)

        self.Plot("Indicators","RSI", self.rsi.Current.Value)

    def OnOrderEvent(self, orderEvent):
        # log order events
        self.Log(str(orderEvent))