#region imports
from AlgorithmImports import *
#endregion
# Your New Python File

def macd(macd, signal):
    m = macd
    s = signal
    if s > 0 and m > s:
        return 1
    elif m >= 0  and m <= s:
        return 2
    elif m < 0  and s >= 0:
        return 3
    elif s < 0  and m <= s:
        return 4
    elif s < m  and m <= 0:
        return 5
    elif s <= 0  and m > 0:
        return 6

#region imports
from AlgorithmImports import *
#endregion
# In Module 5, we address strategy development.
# We will transform informative features into actual investment,
# by making sense of all the observations and formulating a general theory that explains them. 
# We will emphasize the economic mechanism that support the theory, such as behavioral bias, 
# asymmetric information, regulatory constraints. 
# Finally, we code the full algorithm as the prototype



# We will start with a basic option selling strategy
# We sell monthly SPY puts using 90% of our capital without leverage, 
# whenever our portfolio is all cash. We trade at a fixed time during the day.
# We request Options Data Using data.OptionChains instead of OptionsChainProvider
# It returns a list of options contracts for each date.
# There is no indicators because by feature analysis we only know that 
# the probability of a OTM put to expire worthless is pretty high, > 90%.
# If our options get assigned, we sell the stocks next day.

# 11/13/22 Update the algo
# 2/19/21 Complete the barebone algo
# 2/25/21 Adding TA features such as MACD Rating according to WOO system


import pandas as pd

from labels import macd 

class BasicOptionTradingAlgo(QCAlgorithm):

# In Initialize
    def Initialize(self):
        
        # In sample
        self.SetStartDate(2016, 1, 1)
        self.SetEndDate(2021, 1, 1)

        # Out of sample
        # self.SetStartDate(2020, 12, 1)
        # self.SetEndDate(2022, 9, 1)
        
        self.SetCash(400000)
        
        spy = self.AddEquity("SPY", Resolution.Daily)
        equity_symbol = self.AddEquity("SPY", dataNormalizationMode=DataNormalizationMode.Raw).Symbol  # add tsla options
        option = self.AddOption(equity_symbol)
        option.SetFilter(-2, 2, timedelta(0), timedelta(30))
        self.symbol = option.Symbol
        self.SetBenchmark("SPY")
        
        self.underlyings = ['SPY' ]

        # Initialize the dataframe for contracts and positions
        self.len_underlyings = len(self.underlyings)
        self.contracts = [str()]*self.len_underlyings # contract id
        self.shares = [0] * self.len_underlyings # number of shares of underlying assets
        self.label_macd = [0] * self.len_underlyings # MACD rating labels 1-6 by WOO system
        
        
        self.df = pd.DataFrame(index = self.underlyings, data = list(zip(self.contracts, self.shares, self.label_macd)), \
            columns = ['Contract','Number of Shares', 'Label_MACD'])
            
            
        # define our daily macd(12,26) with a 9 day signal 12 is fast peried  26 is slow and 9 is signal
        self.__macd = self.MACD("SPY", 12, 26, 9, MovingAverageType.Exponential, Resolution.Daily)
        # self.__previous = datetime.min
        self.PlotIndicator("MACD", True, self.__macd, self.__macd.Signal)
        # self.PlotIndicator("SPY", self.__macd.Fast, self.__macd.Slow)

        # define our historical volatility with a moving window of 30 days.
        # self.__logr = self.LOGR('SPY', 1, Resolution.Daily)
        # self.logrWin = RollingWindow[IndicatorDataPoint](30)
        # self.__logr.Updated += self.LogrUpdated
        # self.PlotIndicator('LogR', self.__logr)
        
        # Initialize the stock data
        for stock in self.underlyings:
            self.equity = self.AddEquity(stock, Resolution.Minute)
            self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
            
            
        # initialize the option contract container
        self.contract = str()
        self.contractsAdded = set()
        
        # schedule the time to trade
        self.target = "09:46:00"
        self.trade_time = None
        
        # adding position size
        self.frac = 0.9/self.len_underlyings # The frac of the portfolio for exposure; here we use equal weight.
        
    # def LogrUpdated(self, sender, updated):
    #     '''Add updated value to rolling window'''
    #     self.logrWin.Add(updated)
        #self.Log(updated)
        
    def OnData(self, data):
        
        # wait for our macd to fully initialize
        if not self.__macd.IsReady: return
        
        # # wait for logr to fully initialize
        # if not self.LogrWin.IsReady: return
    
        # logr_values = []
        # for i in self.LogrWin:
        #     logr_values.append(i.Value)
        
        
        # hv = np.std(logr_values)*16
    
        # only once per day
        #if self.__previous.date() == self.Time.date(): return

        # define a small tolerance on our checks to avoid bouncing
        # tolerance = 0.0025

        
        Account_Value = self.Portfolio.TotalPortfolioValue
        
        for underlying in self.underlyings:
            '''
            Compute the shares for each underlying. 
            '''
            if not (self.Securities[underlying].Price == 0.0): 
                        # If the option is exerised, the next day's 
                        # stock price data may be zero on the next data slice.
                number_shares = int( round(Account_Value * self.frac /self.Securities[underlying].Price/100) ) * 100
                self.df.at[underlying, 'Number of Shares'] = number_shares 
                self.df.at[underlying, 'Label_MACD'] = macd(self.__macd.Current.Value , self.__macd.Signal.Current.Value) # calculate the label
                # self.df.at[underlying, 'HV'] = np.std(self.__logr)*np.sqrt(252)
                #self.Debug('The number of shares of {} may be put to us is {}.'.format(underlying, number_shares))
                
                # sell the stocks when been put to
                if self.Portfolio[underlying].Invested: 
                    self.MarketOrder(underlying, -self.Portfolio[underlying].Quantity)
                    

        for underlying in self.underlyings:
            '''Update option contracts and trade them'''
            self.contract = self.df.loc[underlying, 'Contract']
            
            if not (self.Securities.ContainsKey(self.contract) and self.Portfolio[self.contract].Invested): 
                
                self.contract = self.OptionsFilter(data)
                self.df.at[underlying, 'Contract'] = str(self.contract) 
                # self.df.at[underlying, 'ImpVol/HisVol'] = self.contract.ImpliedVolatility # *** 
                #self.Debug(str(self.contract))


            if (    self.df.loc[underlying, 'Number of Shares']!= 0 \
                    # To avoid trading zero contract
                    # and self.contract!={} \
                    # To avoid trading an empty contract
                    and self.Securities.ContainsKey(str(self.contract)) and not self.Portfolio[str(self.contract)].Invested \
                    # make sure the time is after the contract is added
                    and self.df.at[underlying, 'Label_MACD'] in [1,2,3,4,5,6] #,3,4,6]
                    # and self.df.at[underlying, 'ImpVol/HisVol'] > 2.0
                    ): 
                    

                    # self.Debug('The contract traded is {}.'.format(contract))
                    
                    self.trade_time = self.Time 
                    if self.trade_time.strftime("%H:%M:%S") == self.target:
                        # Trade not when market just opens since at that moment the price could be zero in the data
                        # https://www.quantconnect.com/forum/discussion/3942/option-examples-from-tutorials-fail/p1/comment-11895
  
                        self.OpenPositions(self.contract.Symbol, underlying)

        
    
    def OpenPositions(self, contract, underlying):
        
        number_contracts = self.df.loc[underlying, 'Number of Shares']/100
        self.MarketOrder(contract, -number_contracts)
        #self.Debug(contract)
 
            
    
    ## Example of a filtering function to be called
    ## The result is a selected option contract
    
    def OptionsFilter(self, data):
        # sort contracts to find at the money (ATM) contract with the farthest expiration
        if data.OptionChains.Count == 0: return
        for i in data.OptionChains:
            if i.Key != self.symbol: continue
            chain = i.Value
            call = [x for x in chain if x.Right == 1] # filter the put options contracts
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(call, key = lambda x: abs(chain.Underlying.Price - x.Strike)),
                                        key = lambda x: x.Expiry, reverse=False)
            if len(contracts) == 0: return
            contract = contracts[0]
            return contract
            

            
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))