# region imports
from AlgorithmImports import *
from datetime import timedelta
# endregion

class LETFShorter(QCAlgorithm):

    def initialize(self):

        #Set up initialization parameters
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_start_date(int(self.get_parameter('Start Year')), int(self.get_parameter('Start Month')), int(self.get_parameter('Start Day')))
        self.set_end_date(int(self.get_parameter('End Year')), int(self.get_parameter('End Month')), int(self.get_parameter('End Day')))
        self.StartingBalance = int(self.get_parameter('Starting Balance'))
        self.set_cash(self.StartingBalance)
        self.set_warm_up(timedelta(1))

        #Set up the benchmark equity for comparison
        self.BMequity = self.add_equity(self.get_parameter('Benchmark Equity'), Resolution.MINUTE)
        self.set_benchmark(self.BMequity.symbol)
        self.BM_priceHistory = []
        self.schedule.on(self.date_rules.every_day(), self.time_rules.before_market_close(self.BMequity.symbol, 0), self.plotBM)

        #Add the LETFs
        self.longLETF = self.add_equity(self.get_parameter('Long 3X LETF'), Resolution.MINUTE)
        self.longLETF.set_data_normalization_mode(DataNormalizationMode.Raw)
        self.shortLETF = self.add_equity(self.get_parameter('Short 3X LETF'), Resolution.MINUTE)
        self.shortLETF.set_data_normalization_mode(DataNormalizationMode.Raw)
        #self.add_option(self.get_parameter('Long 3X LETF'))
        #self.add_option(self.get_parameter('Short 3X LETF'))
        self.longLETF_Ratio = 0
        self.shortLETF_Ratio = 0
        self.longLETF_CALL = None
        self.longLETF_PUT = None
        self.shortLETF_CALL = None
        self.shortLETF_PUT = None
        self.options_Exp = None

        #Add reserve equity
        self.reserveEquity = self.add_equity(self.get_parameter("Reserve Equity"), Resolution.MINUTE)

    def on_data(self, data: Slice):
        #default
        if self.is_warming_up: return

        #liquidate options if they are too close to expiration
        if not self.options_Exp == None:
            if (self.options_Exp - self.time) <= timedelta(int(self.get_parameter('Min Expiration DTE'))):
                self.LiquidateAllOptions()
                #self.log('Closed options positions: too close to expiration.')

        #liquidate options if a split is impending
        if data.Splits.contains_key(str(self.longLETF.symbol)):
            if data.Splits[self.longLETF.symbol].Type == SplitType.Warning:
                self.LiquidateAllOptions()
                self.log('Closed options positions: impending stock split in ' + str(self.longLETF.symbol) + '.')
        elif data.Splits.contains_key(str(self.shortLETF.symbol)):
            if data.Splits[self.shortLETF.symbol].Type == SplitType.Warning:
                self.LiquidateAllOptions()
                self.log('Closed options positions: impending stock split in ' + str(self.shortLETF.symbol) + '.')
        

        if not self.portfolio.invested:
            #self.set_holdings(self.reserveEquity.symbol, float(self.get_parameter("Reserve Allocation")))
            self.BuyContracts(data)


        #Calculate synthetic option position ratios
        self.longLETF_Ratio = (int(self.get_parameter('Long Weighting')) * self.longLETF.price) / ((self.longLETF.price + self.shortLETF.price) * (int(self.get_parameter('Long Weighting')) + int(self.get_parameter('Short Weighting')))) 
        self.shortLETF_Ratio = (int(self.get_parameter('Short Weighting')) * self.shortLETF.price) / ((self.longLETF.price + self.shortLETF.price) * (int(self.get_parameter('Long Weighting')) + int(self.get_parameter('Short Weighting')))) 

    def debugPlot(self):
        return

    def plotBM(self):
        #Plot benchmark on strategy equity graph and plot alpha on alpha graph
        self.BM_priceHistory.append(self.History(self.BMequity.symbol, 2, Resolution.DAILY)['close'].unstack(level= 0).dropna()[self.BMequity.symbol].iloc[-1])
        self.Plot('Strategy Equity', self.BMequity.symbol, self.BM_priceHistory[-1] / self.BM_priceHistory[0] * self.StartingBalance)
        self.Plot('Alpha', 'Alpha', self.portfolio.total_portfolio_value / self.StartingBalance - self.BM_priceHistory[-1] / self.BM_priceHistory[0])

    def BuyContracts(self, data):
        if (self.options_Exp == None):
            self.GetOptions(data)
            return
        elif not (self.Portfolio[self.longLETF_CALL].invested and \
            self.Portfolio[self.longLETF_PUT].invested and \
            self.Portfolio[self.shortLETF_CALL].invested and \
            self.Portfolio[self.shortLETF_PUT].invested):
            if data.contains_key(self.longLETF_CALL) and \
                data.contains_key(self.longLETF_PUT) and \
                data.contains_key(self.shortLETF_CALL) and \
                data.contains_key(self.shortLETF_PUT):
                self.buy(self.longLETF_PUT, 1)
                self.sell(self.longLETF_CALL, 1)
                self.buy(self.shortLETF_PUT, 1)
                self.sell(self.shortLETF_CALL, 1)

    def GetOptions(self, data):
        #get prices of the underlying
        longLETF_underlyingPrice = self.Securities[self.longLETF.symbol].price
        shortLETF_underlyingPrice = self.Securities[self.shortLETF.symbol].price
        #get the list of contracts for each underlying
        longLETF_contracts = self.option_chain_provider.get_option_contract_list(self.longLETF.symbol, data.time)
        shortLETF_contracts = self.option_chain_provider.get_option_contract_list(self.shortLETF.symbol, data.time)
        #first filter the options for calls with an OTM strike greater than the parameter setting and at least the parameter setting for opening DTE
        longLETF_calls = [i for i in longLETF_contracts if i.ID.OptionRight == OptionRight.CALL and \
            i.ID.StrikePrice - longLETF_underlyingPrice > (1 + float(self.get_parameter('OTM Percentage'))) * longLETF_underlyingPrice and \
                int(self.get_parameter('Open Expiration DTE')) < (i.ID.Date - data.time).days]
        shortLETF_calls = [i for i in shortLETF_contracts if i.ID.OptionRight == OptionRight.CALL and \
            i.ID.StrikePrice - shortLETF_underlyingPrice > (1 + float(self.get_parameter('OTM Percentage'))) * shortLETF_underlyingPrice and \
                int(self.get_parameter('Open Expiration DTE')) < (i.ID.Date - data.time).days]
        #temporarily store old expiration value
        old_exp = self.options_Exp
        self.options_Exp = None
        #now we need to find a common expiration date between both option chains for both underlyings
        if len(longLETF_calls) > 0 and len(shortLETF_calls) > 0:
            for i in sorted(longLETF_calls, key = lambda x: abs((x.ID.Date - self.Time).days)):
                for z in sorted(shortLETF_calls, key = lambda x: abs((x.ID.Date - self.Time).days)):
                    if i.ID.Date == z.ID.Date:
                        self.options_Exp = i.ID.Date
                        break
                else:
                    continue
                break
            if not self.options_Exp == None:
                #now we sort the calls to find the appropriate call to sell for each underlying, using the contract date and then sort again based on the strike price
                longLETF_call_contract = sorted(sorted(longLETF_calls, key = lambda x: abs((x.ID.Date - self.Time).days - (self.options_Exp - self.Time).days)), key = lambda x: x.ID.StrikePrice - longLETF_underlyingPrice)[0]
                shortLETF_call_contract = sorted(sorted(shortLETF_calls, key = lambda x: abs((x.ID.Date - self.Time).days - (self.options_Exp - self.Time).days)), key = lambda x: x.ID.StrikePrice - shortLETF_underlyingPrice)[0]
                #now that we have the calls, we have to get puts with matching strikes and expirations.  This is easier because we did most of the work to select the call.
                longLETF_puts = [i for i in longLETF_contracts if i.ID.OptionRight == OptionRight.PUT and i.ID.StrikePrice == longLETF_call_contract.ID.StrikePrice and i.ID.Date == self.options_Exp]
                shortLETF_puts = [i for i in shortLETF_contracts if i.ID.OptionRight == OptionRight.PUT and i.ID.StrikePrice == shortLETF_call_contract.ID.StrikePrice and i.ID.Date == self.options_Exp]
                if len(longLETF_puts) > 0 and len(shortLETF_puts) > 0:
                    self.longLETF_CALL = longLETF_call_contract
                    self.add_option_contract(longLETF_call_contract, Resolution.MINUTE)
                    self.longLETF_PUT = longLETF_puts[0]
                    self.add_option_contract(longLETF_puts[0], Resolution.MINUTE)
                    self.shortLETF_CALL = shortLETF_call_contract
                    self.add_option_contract(shortLETF_call_contract, Resolution.MINUTE)
                    self.shortLETF_PUT = shortLETF_puts[0]
                    self.add_option_contract(shortLETF_puts[0], Resolution.MINUTE)
            else:
                self.options_Exp = old_exp
                
    def LiquidateAllOptions(self):
        if not self.longLETF_CALL == None:
            self.remove_option_contract(self.longLETF_CALL)
        if not self.longLETF_PUT == None:
            self.remove_option_contract(self.longLETF_PUT)
        if not self.shortLETF_CALL == None:
            self.remove_option_contract(self.shortLETF_CALL)
        if not self.shortLETF_PUT == None:
            self.remove_option_contract(self.shortLETF_PUT)
        self.longLETF_CALL = None
        self.longLETF_PUT = None
        self.shortLETF_CALL = None
        self.shortLETF_PUT = None
        self.options_Exp = None

    def OnOrderEvent(self, orderEvent):
        self.log(str(orderEvent))
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if order.Type == OrderType.OptionExercise:
            self.Liquidate(orderEvent.Symbol.Underlying)
            self.LiquidateAllOptions()