from typing import Dict, List

# 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 = True  # 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, 8, 5)
        self.SetCash(10000000)

        future = self.AddFuture(Futures.Indices.NASDAQ100EMini, Resolution.Minute) 
        future.SetFilter(0, 90)
        
        self.rsi = RelativeStrengthIndex(RSI_period)

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

        self.just_bought = False
        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

        self.prev_contract = None


        self.stoploss_cumulative = 0
        self.takeprofit_cumulative = 0

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

        if not self.Securities[self.GetSymbol()].IsTradable:
            return
            
    def GetSymbol(self) -> Symbol:
        '''
        get current front month contract
        '''
        return list(self.consolidators.keys())[0]

    def NewDay(self, data:Slice):
        self.Plot('Stoploss count', 'count', self.stoploss_cumulative)
        self.Plot('Take profit count', 'count', self.takeprofit_cumulative)

        self.Reset()
        self.PruneContracts()
        self.ChooseContract(data)
    
    def Reset(self):
        self.Liquidate()
        self.last_rsi = None
        self.loss_count = 0
        self.rsi.Reset()
        self.ResetOrders()

    def ChooseContract(self, data:Slice):
        for chain in data.FutureChains:
            contracts = [contract for contract in chain.Value]
            if len(contracts) == 0:
                self.Print('No contracts found')
                return
            
            contract = sorted(contracts, key=lambda k : k.OpenInterest, reverse=True)[0]
            self.SetUpContract(contract.Symbol)
            continue

    def SetUpContract(self, symbol:Symbol):
        # consolidate the contract data
        consolidator = QuoteBarConsolidator(bar_size)
        consolidator.DataConsolidated += self.OnDataConsolidated
        self.SubscriptionManager.AddConsolidator(symbol, consolidator)
        self.consolidators[symbol] = consolidator

    def PruneContracts(self):
        # unconsolidate untradeable symbols
        for symbol in self.consolidators:
            consolidator = self.consolidators[symbol]
            self.SubscriptionManager.RemoveConsolidator(symbol, consolidator)
            consolidator.DataConsolidated -= self.OnDataConsolidated
        self.consolidators = {}

    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.just_bought = True
                self.market_ticket = self.MarketOrder(symbol, quantity)
                
        if quoteBar.High > self.high:
            self.high = quoteBar.High 
            if self.stoploss_ticket:
                updateFields = UpdateOrderFields()
                updateFields.StopPrice = self.high - stoploss_dist
                updateFields.LimitPrice = self.high - stoploss_dist - stoplimit_dist
                self.stoploss_ticket.Update(updateFields) 
        if self.takeprofit and quoteBar.Close > self.takeprofit:
            self.takeprofit_cumulative += 1
            self.Liquidate()
            self.ResetOrders()

        self.last_rsi = curr_rsi

    def OnOrderEvent(self, orderEvent: OrderEvent):
        if self.just_bought:
            self.just_bought = False
            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
            self.stoploss_cumulative += 1