from typing import Dict

# CONFIGS
RSI_period = 14
RSI_upper = 30
RSI_lower = 27
bar_size = timedelta(minutes=5)

portfolio_pct = .1
stoploss_dist = 20 # distance below high for stop loss
takeprofit_dist = 20 # distance above high for take profit
stoplimit_dist = 5 # distance b/w the stop and limit prices of a stop limit order
max_losses = 2 # max number of losses in a day
debug = False  # turn off to reduce logging
# END CONFIGS

f = False
if f:
    from AlgorithmImports import *

class Consulting(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2019, 8, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(10000000)

        future = self.AddFuture(Futures.Indices.NASDAQ100EMini, Resolution.Minute) 
        future.SetFilter(lambda x: x.FrontMonth().OnlyApplyFilterAtMarketOpen())
        
        self.rsi = RelativeStrengthIndex(RSI_period)

        self.consolidators: Dict[Symbol, QuoteBarConsolidator] = {}

        self.market_ticket = None # Market order ticket
        self.stoploss_ticket = None  # StopMarket stop loss ticket
        self.takeprofit = None # take profit price
        self.high = -1
        self.last_rsi = None

        self.loss_count = 0
        self.curr_day = -1

        self.quoteBar = None
        
    def OnData(self, data:Slice):
        # new day
        if self.curr_day != self.Time.day:
            self.curr_day = self.Time.day 
            self.Reset()
            return

        if not self.Securities[self.GetSymbol()].IsTradable:
            return

        symbol = self.GetSymbol()

        price = self.Securities[symbol].Price
        
        if data.QuoteBars.ContainsKey(symbol):
            high = data[symbol].High
            if high > self.high:
                self.high = high
                # update stop loss if new high reached 
                if self.stoploss_ticket:
                    updateFields = UpdateOrderFields()
                    updateFields.StopPrice = high - stoploss_dist
                    updateFields.LimitPrice = high - stoploss_dist - stoplimit_dist
                    self.stoploss_ticket.Update(updateFields) 
            elif self.takeprofit and price > self.takeprofit:
                self.Liquidate()
                self.ResetOrders()
        
            
        
    def GetSymbol(self) -> Symbol:
        '''
        get current front month contract
        '''
        # if len(self.consolidators) > 0:
        #     self.Print('more than one contract')
        return list(self.consolidators.keys())[0]

    def Reset(self):
        self.last_rsi = None
        self.loss_count = 0

        self.rsi.Reset()

    def ResetOrders(self):
        self.market_ticket = None 
        self.stoploss_ticket = None 
        self.takeprofit = None

    def Print(self, msg):
        if debug:
            self.Log(msg)

    def OnDataConsolidated(self, sender, quoteBar:QuoteBar):
        '''
        5 minute consolidator
        Update RSI, SetHoldings
        '''

        self.rsi.Update(self.Time, quoteBar.Close)
        
        self.quoteBar = quoteBar

        if not self.rsi.IsReady:
            return 
        
        curr_rsi = self.rsi.Current.Value
        
        if self.loss_count > max_losses:
            return

        # self.Plot('RSI', 'Value', curr_rsi)

        if (not self.Portfolio.Invested and self.last_rsi 
                and self.last_rsi < RSI_lower and curr_rsi > RSI_upper):
            if quoteBar.High - stoploss_dist < quoteBar.Close:
                symbol = self.GetSymbol()
                quantity = self.CalculateOrderQuantity(symbol, portfolio_pct)
                self.market_ticket = self.MarketOrder(symbol, quantity)
                
        self.last_rsi = curr_rsi

    def OnOrderEvent(self, orderEvent: OrderEvent):
        if self.market_ticket and orderEvent.OrderId == self.market_ticket.OrderId:
            self.stoploss_ticket = self.StopLimitOrder(orderEvent.Symbol, -orderEvent.Quantity, self.quoteBar.High - stoploss_dist, self.quoteBar.High - stoploss_dist - stoplimit_dist) 
            self.takeprofit = self.quoteBar.High + stoploss_dist
        elif self.stoploss_ticket and orderEvent.OrderId == self.stoploss_ticket.OrderId:
            self.ResetOrders()
            self.loss_count += 1
        
    def OnSecuritiesChanged(self, changes):
        '''
        register consolidator for frontmonth contract, deregister consolidator for expired contract
        '''
        self.Reset()

        self.Print('changing securities')
        new_symbol = None

        for security in changes.AddedSecurities:
            new_symbol = security.Symbol
            consolidator = QuoteBarConsolidator(bar_size)
            consolidator.DataConsolidated += self.OnDataConsolidated
            self.SubscriptionManager.AddConsolidator(new_symbol, consolidator)
            self.consolidators[new_symbol] = consolidator
            continue
        
        # remove old symbols
        if new_symbol:
            for symbol in list(self.consolidators.keys()):
                if new_symbol == symbol:
                    continue
                self.Print('removing old contract')
                consolidator = self.consolidators.pop(symbol)
                self.SubscriptionManager.RemoveConsolidator(symbol, consolidator)
                consolidator.DataConsolidated -= self.OnDataConsolidated