import numpy as np
import pandas as pd
from collections import deque

class QuantumOptimizedPrism(QCAlgorithm):

    def Initialize(self):
        
        # Max window size
        self.Lookback = 168 # 1 week in hours
    
        # list of symbols we want to trade
        self.symbolList = ["BTCUSD","ETHUSD","LTCUSD","BCHUSD"]
        #self.symbolList = ["BTCUSD","ETHUSD","LTCUSD","BCHUSD","XRPUSD","XLMUSD","EOSUSD","REPUSD","XTZUSD","ETCUSD","ZRXUSD"]
        
        # dictionary to store info per symbol
        self.rollingWindow = {}
        
        # init insights for plotting alpha
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(Resolution.Hour), self.Lookback)
        self.insightTemp = None
        
        # create plot and consolidator for each symbol we want to trade
        for name in self.symbolList:
            cryptoSymbol = self.AddCrypto(name, Resolution.Minute, Market.GDAX).Symbol
            self.Debug('{}'.format(cryptoSymbol))
            fiveConsolidator = TradeBarConsolidator(timedelta(hours=1))
            fiveConsolidator.DataConsolidated += self.FiveConsolidator
            self.SubscriptionManager.AddConsolidator(cryptoSymbol, fiveConsolidator)
            self.rollingWindow["close_{0}".format(cryptoSymbol)] = deque(maxlen=self.Lookback)
            
            stockPlot = Chart(str(cryptoSymbol))
            stockPlot.AddSeries(Series('Price', SeriesType.Line, 0))
            stockPlot.AddSeries(Series('Top', SeriesType.Line, 0))
            stockPlot.AddSeries(Series('Mid', SeriesType.Line, 0))
            stockPlot.AddSeries(Series('Bot', SeriesType.Line, 0))
            self.AddChart(stockPlot)
        
        # max buy sell amount
        self.weight = 1 / len(self.symbolList)
        
        self.SetCash(100000)
        self.SetStartDate(2020, 1, 1)
        self.SetBrokerageModel(BrokerageName.GDAX, AccountType.Cash)
        self.SetWarmUp(self.Lookback)

    def GetInsight(self, bar, price, top, mid, bot):
        '''Updates this alpha model with the latest data from the consolidator'''
        
        insights = []
        if price > top:
            direction = InsightDirection.Up
        elif price < top and price > mid:
            direction = InsightDirection.Flat
        else:
            direction = InsightDirection.Down
        
        if self.insightTemp is not None and direction == self.insightTemp:
            return insights
        
        self.insightTemp = direction
            
        insight = Insight.Price(bar.Symbol, self.insightPeriod, direction)
        insights.append(insight)
        return insights
        
    def Indicator(self, symbol):
        '''Rolling quantile for upper and lower bounds'''
        data = np.array(self.rollingWindow["close_"+str(symbol)])
        top = np.quantile(data,0.99, interpolation='higher')
        bot = np.quantile(data,0.01, interpolation='higher')
        mid = (top+bot)/2
        return top, bot, mid
            
    def FiveConsolidator(self, sender, bar):
        
        stringSym = str(bar.Symbol)
        sym_price = bar.Close
        
        # Append price data
        self.rollingWindow["close_{0}".format(bar.Symbol)].append(sym_price)
        
        if not self.IsWarmingUp:
            if len(self.rollingWindow["close_{0}".format(bar.Symbol)]) < self.Lookback:
                return
            
            # Get indicators
            top, bot, mid = self.Indicator(bar.Symbol)
            
            # buy
            if sym_price > top and not self.Portfolio.Invested:
                self.SetHoldings(bar.Symbol, self.weight)
                
            # sell
            elif sym_price < mid and self.Portfolio.Invested:
                self.SetHoldings(bar.Symbol, 0)
                
            else:
                pass
            
            # plot insights (not needed for trading).  Disable for faster backtesting
            self.EmitInsights(self.GetInsight(bar, sym_price, top, mid, bot))
            
            # plot every interval defined plot!
            if self.Time.hour % 24 == 0 and self.Time.minute == 0:
                self.Plot(str(bar.Symbol), 'Price', sym_price)
                self.Plot(str(bar.Symbol), 'Top', top)
                self.Plot(str(bar.Symbol), 'Mid', mid)
                self.Plot(str(bar.Symbol), 'Bot', bot)

        
    def OnData(self, data):
        pass

    def OnOrderEvent(self, orderEvent):
        self.Debug("{} {}".format(self.Time, orderEvent.ToString()))

    def OnEndOfAlgorithm(self):
        self.Log("{} - TotalPortfolioValue: {}".format(self.Time, self.Portfolio.TotalPortfolioValue))
        self.Log("{} - CashBook: {}".format(self.Time, self.Portfolio.CashBook))