import clr
clr.AddReference("System")
clr.AddReference("QuantConnect.Algorithm")
clr.AddReference("QuantConnect.Common")
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
#from QuantConnect.Algorithm.Framework import *
#from QuantConnect.Algorithm.Framework.Alphas import *
#from QuantConnect.Algorithm.Framework.Execution import *
#from QuantConnect.Algorithm.Framework.Risk import *
#from QuantConnect.Algorithm.Framework.Selection import *
#from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Indicators import *
from QuantConnect.Data.Market import TradeBar
from QuantConnect.Python import *
from QuantConnect.Storage import *
import numpy as np
from QuantConnect import *
from QuantConnect.Parameters import *
from QuantConnect.Benchmarks import *
from QuantConnect.Brokerages import *
from QuantConnect.Util import *
from QuantConnect.Interfaces import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Indicators import *
from QuantConnect.Data import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Data.Custom import *
from QuantConnect.Data.Fundamental import *
from QuantConnect.Data.Market import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Notifications import *
from QuantConnect.Orders import *
from QuantConnect.Orders.Fees import *
from QuantConnect.Orders.Fills import *
from QuantConnect.Orders.Slippage import *
from QuantConnect.Scheduling import *
from QuantConnect.Securities import *
from QuantConnect.Securities.Equity import *
from QuantConnect.Securities.Forex import *
from QuantConnect.Securities.Interfaces import *
from datetime import date, datetime, timedelta
from QuantConnect.Python import *
from QuantConnect.Storage import *
from datetime import datetime
import decimal
import heapq
from itertools import groupby
from heapq import heappush
import random
import math
from decimal import *
from QuantConnect import Market
class orderimbalance(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 10, 8)
        self.SetEndDate(2015, 10, 8)
        self.SetCash(1000000)
        self.symbol=self.AddEquity("SPY", Resolution.Tick, Market.USA, True, 0, True).Symbol
        
        self.Nlevel = 5
        self.W = [math.exp( -i ) for i in range(self.Nlevel)]
        self.cont = 0
        self.Lb = []
        self.Ls = []
        
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.Every(timedelta(minutes=5)), self.plotting)
    def OnData(self, data):
        #if not data.ContainsKey(self.symbol) and not  data.Ticks.ContainsKey(self.symbol):
        #    return
        if  self.cont<300000:
            numTick=0
            for tick in data.Ticks[self.symbol]:
                numTick+=1
                if tick.TickType == TickType.Trade:
                    if self.Lb and self.Ls :
                        if tick.Price in [-x for x in self.Lb[:][0]]:
                            p, v, d = tick.Price, tick.Quantity, 1
                            order=[[p,v,d]]
                            self.Lb, self.Ls = self.getLOrders(order, self.Lb, self.Ls)
                        if tick.Price in self.Ls[:][0]:
                            p, v, d = tick.Price, tick.Quantity, 0
                            order=[[p,v,d]]
                            self.Lb, self.Ls = self.getLOrders(order, self.Lb, self.Ls)

                if tick.TickType == TickType.Quote:
                    if tick.AskPrice>1:
                        p, v, d = tick.AskPrice, tick.AskSize, 1
                    else:
                        p, v, d = tick.BidPrice, tick.BidSize, 0

                    order=[[p,v,d]]
                    self.Lb, self.Ls = self.getLOrders(order, self.Lb, self.Ls)
            self.cont+=1

    def plotting(self):
        if self.Lb and self.Ls and self.Time.hour < 14 and self.Time.hour >8:
            self.Debug(self.Time)
            
            self.Plot("Price", "Bid2", -self.Lb[1][0])
            self.Plot("Price", "Bid1", -self.Lb[0][0])
            self.Plot("Price", "Ask1", self.Ls[0][0])
            self.Plot("Price", "Ask2", self.Ls[1][0])
            self.Plot("Price", "MidPrice", (-self.Lb[0][0]+self.Ls[0][0])/2)
            
            Lb = [[float(Decimal("%.2f"%(-item[0]))) , int(Decimal("%.2f"%(item[1])))] for item in self.Lb]
            Ls = [[float(Decimal("%.2f"%(item[0]))) , int(Decimal("%.2f"%(item[1])))] for item in self.Ls]
            sumVolBid = sum(Vb*W for Vb,W in zip(Lb[0:self.Nlevel][1],self.W))
            sumVolAsk = sum(Va*W for Va,W in zip(Ls[0:self.Nlevel][1],self.W))
            imbalanceRatio = (sumVolBid-sumVolAsk)/(sumVolBid+sumVolAsk)
            self.Plot("Imbalance Ratio", "Value", imbalanceRatio)
            
            self.Debug(str(Lb[self.Nlevel-1::-1])+"   "+str(Ls[0:self.Nlevel]))
        if self.Time.hour > 14:
            self.Quit()
    
    def getLOrders(self,orders,bu,se):
        for x in orders:
            
            a,b,t = x[0],x[1],x[2]
            n = b
            if t == 0: # buy
                while n > 0 and se:
                    p,num = heapq.heappop(se) 
                    if p > a:
                        heapq.heappush(se, [p, num]) 
                        break
                    if num <= n:
                        n -= num
                        while n>0:
                            p,num = heapq.heappop(se)
                            if num <= n:
                                n -= num
                            else:
                                heapq.heappush(se, [p, num-n])
                                n = 0
                    else: 
                        heapq.heappush(se, [p, num-n]) 
                        n = 0 
                if n > 0: 
                    heapq.heappush(bu,[-a,n])
            else:
                while n > 0 and bu:
                    p, num = heapq.heappop(bu)
                    if -p < a:
                        heapq.heappush(bu, [p, num])
                        break
                    if num <= n:
                        n -= num
                        while n>0:
                            p,num = heapq.heappop(bu)
                            if num <= n:
                                n -= num
                            else:
                                heapq.heappush(bu, [p, num-n]) # fill all that order then we have num-n volume
                                n = 0
                    else:
                        heapq.heappush(bu, [p, num-n])
                        n = 0
                if n > 0:
                    heapq.heappush(se,[a,n])
        bu=[[k, sum(v for _, v in g)] for k, g in groupby(sorted(bu), key = lambda x: x[0])]
        se=[[k, sum(v for _, v in g)] for k, g in groupby(sorted(se), key = lambda x: x[0])]
        return bu, se