from QuantConnect.Indicators import *

class CryptoArb(QCAlgorithm):

    # ==================================================================================
    # Main entry point for the algo     
    # ==================================================================================
    def Initialize(self):
        self.InitAlgoParams()
        #self.Schedule.On(self.DateRules.EveryDay(),
        #         self.TimeRules.Every(timedelta(minutes=1)),
        #         self.Arbitrage)

    # ==================================================================================
    # Set algo params: Symbol, broker model, ticker, etc. Called from Initialize(). 
    # ==================================================================================
    def InitAlgoParams(self):
        
        # Set params for data and brokerage
        # -----------------------------------
        self.SetBrokerageModel(BrokerageName.Bitfinex, AccountType.Margin)
        self.initCash = 20000           # todo: use this to track buy+hold     
        self.SetStartDate(2018, 9, 10)  # Set Start Date 
        self.SetEndDate (2018, 9, 11) # Set End Date
        self.SetCash(self.initCash)     # Set Strategy Cash
        self.ETHUSD = self.AddCrypto("ETHUSD", Resolution.Tick).Symbol # ETH in USD
        self.BTCUSD = self.AddCrypto("BTCUSD", Resolution.Tick).Symbol # BTC in USD
        self.ETHBTC = self.AddCrypto("ETHBTC", Resolution.Tick).Symbol # ETH in BTC
        self.triangle = (
            self.ETHUSD,
            self.BTCUSD,
            self.ETHBTC
        )
        self.SetWarmup(10)
        # Tracking variables
        self.imbalance = None
        self.next_index = None
        self.buy_or_sell = 0 # none = 0, buy = 1, sell = 2
        self.trade_lock = False
        # Limit trades
        self.max_iterations = 20
        self.iteration = 0
        
    def find_imbalance(self, triangle):
        """Returns the index of the crypto pair which is out of balance.
        if the first crypto pair is cheaper than its synthetic, we want to arbitrage it
        if the first crypto pair is more expensive than its synthetic, we want to arbitrage the second pair
        """
        sym_1 = triangle[0] # BTCUSD
        sym_1_price = self.Securities[sym_1].Price
        sym_2 = triangle[1] # ETHUSD
        sym_2_price = self.Securities[sym_2].Price
        sym_3 = triangle[2] # BTCETH
        sym_3_price = self.Securities[sym_3].Price
        # Check imbalance (as percentage of asset price)
        imbalance = (sym_1_price - (sym_2_price*sym_3_price)) / sym_1_price # E.g. USD/ETH - USD/BTC*BTC/ETH
        return imbalance
    
    def reset_variables(self):
        self.imbalance = None
        self.next_index = None
        self.buy_or_sell = 0
        self.trade_lock = False
        self.iteration = self.iteration + 1
    
    def order_symbol_allocation(self, symbol, allocation):
        quantity = self.CalculateOrderQuantity(symbol, allocation)
        return self.MarketOrder(symbol, quantity)

    # ========================================================================
    # Plot Charts. Called whenever we want to plot current values.
    # ========================================================================
    def PlotCharts(self):
        self.Plot('Prices', 'ETHUSD', self.Securities[self.ETHUSD].Price)
        self.Plot('Prices', 'BTCUSD', self.Securities[self.BTCUSD].Price)
        self.Plot('Prices', 'ETHBTC', self.Securities[self.ETHBTC].Price)
    

    # ========================================================================
    # State change logic for handling order flow
    # ========================================================================
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        # We only want to execute when an order has been filled
        if order.Status != OrderStatus.Filled: 
            return
        # Second stage of the arbitrage
        ## Handle the intermediate step (exchanging the two cryptos for each other)
        if self.next_index == 2:
            # Intermediate sell
            if self.buy_or_sell == 2:
                self.Debug("Step 2: Selling {}".format(self.triangle[self.next_index]))
                self.next_index = 1 # Finally sell the alternate crypto for USD
                self.buy_or_sell = 2 # Need to sell back to USD
                self.SetHoldings(self.triangle[self.next_index], 0.0) # Sell the primary crypto for the alternate crypto
                return
            # Intermediate buy
            elif self.buy_or_sell == 1:
                self.Debug("Step 2: Buying {}".format(self.triangle[self.next_index]))
                self.next_index = 0 # Finally sell the primary crypto for USD
                self.buy_or_sell = 2 # Need to sell back to USD
                self.SetHoldings(self.triangle[self.next_index], 1.0) # Buy the primary crypto with the alternate crypto
                return
        # Final stage of the arbitrage
        ## Handle the final step (exchanging the crypto back to USD)
        elif self.next_index < 2 and self.next_index >= 0:
            # Final sell
            self.Debug("Step 3: Selling {}".format(self.triangle[self.next_index]))
            self.next_index = -1 # Last trade
            self.SetHoldings(self.triangle[self.next_index], 0.0) # Sell the primary or alternate crypto for USD
            return
        # Exited
        elif self.next_index == -1:
            self.reset_variables()
                    
    # ============================================================================
    # OnData handler. Triggered by data event (i.e. every time there is new data).
    # ============================================================================ 
    # todo: track 'close' on multiple resolutions (daily bars and hourly bars)
    def OnData(self, data):
        if self.IsWarmingUp:
            return
    
        # Exit on max iterations
        if self.iteration > self.max_iterations:
            self.Liquidate()
            self.Quit()
        
        # Don't initiate if we are in the middle of an arbitrage trade
        if self.trade_lock:
            return
    
        # We are not currently in the middle of an arbitrage trade
        if not self.imbalance and self.next_index is None:# or abs(self.imbalance) < 1:
            self.imbalance = self.find_imbalance(self.triangle)
    
        # TODO: Calculate if trade will be worth it after transaction fees
        if abs(self.imbalance) < .01:
            self.imbalance = None
            return
            
        # First stage of the arbitrage
        ## Initiates the order cascade
        if self.next_index is None and not self.trade_lock:
            self.Debug("imbalance: {}, next_index: {}, buy_or_sell: {}".format(self.imbalance, self.next_index, self.buy_or_sell))
            # We are going to start an arbitrage trade, lock trading
            self.trade_lock = True
            ## If the imbalance is negative, we buy the primary crypto
            if self.imbalance < 0:
                self.Debug("Step 1: Buying {}".format(self.triangle[0]))
                self.next_index = 2 # Need to sell to get the intermediate crypto next
                self.buy_or_sell = 2 # Need to sell the intermediate next
                self.SetHoldings(self.triangle[0], 1.0) # Buy the primary crypto
                return
            elif self.imbalance > 0:
                self.Debug("Step 1: Buying {}".format(self.triangle[1]))
                self.next_index = 2 # Need to buy the primary crypto with the alternate
                self.buy_or_sell = 1 # Need to buy the intermediate next
                self.SetHoldings(self.triangle[1], 1.0) # Buy the alternate crypto
                return