from math import floor


class JumpingVioletDonkey(QCAlgorithm):

    def Initialize(self):

        # model parameters
        self.stopInitialThreshold = 0.98
        self.limitThreshold = 0.99
        self.instrument = "TSLA"
        self.momentumPeriod = 4
        self.minimumMomentumPercentage = 0.01
        self.instrumentResolution = Resolution.Second

        # backtesting parameters
        self.initialCash = 100000
        self.invPercent = 0.2
        self.SetStartDate(2021, 2, 16)  # Set Start Date
        self.SetEndDate(2021, 2, 23)  # Set Stop Date
        self.SetCash(self.initialCash)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage) 
        self.SetWarmUp(20)

        # setup equity
        self.s = self.AddEquity(self.instrument, self.instrumentResolution)
        self.s.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.SetBenchmark(SecurityType.Equity, self.instrument)
        
        self.mom = self.MOMP(self.instrument, self.momentumPeriod)
        
        self.Schedule.On(self.DateRules.EveryDay(self.instrument), self.TimeRules.At(9,15),  self.OpenPositions)
        self.Schedule.On(self.DateRules.EveryDay(self.instrument), self.TimeRules.At(13,45), self.ClosePositions)

        # state initialization
        self.ticketPostMarkedList = {}
        self.lastGenOrderId = 0
        self.tradingOpen = False

        self.ResetTickets()



    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        if not self.tradingOpen:
            return
        
        
        # self.Plot("momp", self.mom)
        # self.Debug("  - mom = " + str(self.mom.Current.Value) + " @ " + str(self.Time))
            
        if not self.Portfolio.Invested and self.ticketBuy is None:
            
            if not self.mom.Current.Value > self.minimumMomentumPercentage:
                return
            
            self.numSharesBuy = self.CalculateOrderQuantity(self.instrument, self.invPercent)
            self.initialBar = self.Securities[self.instrument]
            self.Debug("- initial bar: " + self.Bar2Str(self.initialBar))
            self.initialClose = self.Securities[self.instrument].Close
            self.limitPrice = round(self.initialClose * self.limitThreshold, 2)
            self.ticketBuyLimitTag = self.GenerateTag("buyLimit")
            self.ticketBuy = self.LimitOrder(self.instrument, self.numSharesBuy, self.limitPrice, self.ticketBuyLimitTag)
            self.Debug("- Submitted limit ticketBuy @ " + str(self.limitPrice))

    def OnOrderEvent(self, orderEvent):

        ticket = self.Transactions.GetOrderTicket(orderEvent.OrderId)
        self.Debug("- OnOrderEvent: {0} {1} {2}".format(ticket.Tag, orderEvent.Status, ticket))
        
        if self.IsPostFillReady(ticket, self.ticketBuyLimitTag):
            self.Debug("- Filled ticketBuy @ " + str(orderEvent.FillPrice))
            self.stopPrice = round(self.initialClose * self.stopInitialThreshold, 2)
            self.ticketSellStopTag = self.GenerateTag("sellStop")
            self.ticketStop = self.StopMarketOrder(self.instrument, - self.numSharesBuy, self.stopPrice, self.ticketSellStopTag)
            self.Debug("- Submitted stop ticketStop @ " + str(self.stopPrice))
            self.MarkPostFillDone(ticket)

        if self.IsPostFillReady(ticket, self.ticketSellStopTag):
            self.Debug("- Filled ticketStop @ " + str(orderEvent.FillPrice))
            self.ResetTickets()
            # TODO - add cool-off period
            self.MarkPostFillDone(ticket)

    def ResetTickets(self):
        self.numSharesBuy = 0

        self.initialClose = 0
        self.limitPrice = 0
        self.ticketBuyLimitTag = ""
        self.ticketBuy = None

        self.stopPrice = 0
        self.ticketSellStopTag = ""
        self.ticketStop = None

        self.ticketStop = None
        
    def Bar2Str(self, bar):
        return "HOLC={0}/{1}/{2}/{3}".format(bar.High, bar.Open, bar.Close, bar.Low)

    def GenerateTag(self, orderName):
        self.lastGenOrderId += 1
        return "tag" + orderName + str(self.lastGenOrderId)

    def OpenPositions(self):
        self.tradingOpen = True
        self.Debug("- Market opened @ " + str(self.Time))

    def ClosePositions(self):
        self.tradingOpen = False
        self.Liquidate("TSLA")
        self.Debug("- Market closed @ " + str(self.Time))
    

    # utility function to test if a ticket order is filled and ticked not post-marked yet 
    def IsPostFillReady(self, ticket, expectedTicketTag):
        filter = ticket is not None and ticket.Tag == expectedTicketTag and ticket.Status == OrderStatus.Filled
        if not filter:
            return False
        if ticket.OrderId in self.ticketPostMarkedList.keys():
            return False
        return True

    # utility function to mark a ticked as post-filled (i.e. post-filled processing was done)
    def MarkPostFillDone(self, ticket):
        self.ticketPostMarkedList[ticket.OrderId] = True