from QuantConnect.Data.Market import TradeBar
from datetime import timedelta
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
import decimal as d
from statistics import pstdev
from System.Drawing import Color

class VWAPAlgorithmicTrader(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 2, 28)
        self.SetEndDate(2022, 3, 7)
        self.SetCash(25000)
        self.ticker = "QQQ"
        self.consPeriod = 1
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        
        #rolling window (to be implemented later)
        self.window = RollingWindow[TradeBar](2)
        
        #symbol and consolidator
        self.symbol = self.AddEquity(self.ticker, Resolution.Minute).Symbol
        self.Consolidate(self.symbol, timedelta(minutes=self.consPeriod), self.BarHandler)
        
        #vwap indicator, consolidator, and vwap_values list for calculating std dev
        self.vwap = self.VWAP(self.symbol, 390)
        self.RegisterIndicator(self.symbol, self.vwap, timedelta(minutes=self.consPeriod))
        self.vwap_values = []
        
        #vwap standard deviation indicator
        self.vwap_stddev = StandardDeviation(390)
        self.RegisterIndicator(self.symbol, self.vwap_stddev, timedelta(minutes=self.consPeriod))
        
        self.vwap.Updated += self.OnVWAPUpdated
        
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.BeforeMarketClose(self.symbol, 0), self.VWAPReset)
        
        #charting logic
        stockPlot = Chart('Indicator Chart')
        
        stockPlot.AddSeries(Series('Price', SeriesType.Line, '', Color.Lime))
        stockPlot.AddSeries(Series('STD T2', SeriesType.Line, '', Color.WhiteSmoke))
        stockPlot.AddSeries(Series('STD T1', SeriesType.Line, '', Color.WhiteSmoke))
        stockPlot.AddSeries(Series('VWAP', SeriesType.Line, '', Color.DarkOrchid))
        stockPlot.AddSeries(Series('STD B1', SeriesType.Line, '', Color.WhiteSmoke))
        stockPlot.AddSeries(Series('STD B2', SeriesType.Line, '', Color.WhiteSmoke))
        self.AddChart(stockPlot)
        
        
    def OnData(self, data):
        
        self.window.Add(data[self.ticker])

        if not self.window.IsReady:
            return
        
        self.STDT2 = self.vwap.Current.Value + (2 * self.vwap_stddev.Current.Value)
        self.STDT1 = self.vwap.Current.Value + (1 * self.vwap_stddev.Current.Value)
        self.STDB1 = self.vwap.Current.Value - (1 * self.vwap_stddev.Current.Value)
        self.STDB2 = self.vwap.Current.Value - (2 * self.vwap_stddev.Current.Value)
        
        self.Plot('Indicator Chart', 'STD T2', self.STDT2)
        self.Plot('Indicator Chart', 'STD T1', self.STDT1)
        self.Plot('Indicator Chart', 'VWAP', self.vwap.Current.Value)
        self.Plot('Indicator Chart', 'STD B1', self.STDB1)
        self.Plot('Indicator Chart', 'STD B2', self.STDB2)
        self.Plot('Indicator Chart', 'Price', data.Bars[self.symbol].Close)
        

    def OnVWAPUpdated(self, vwap, current):
        if vwap.IsReady:
            self.vwap_stddev.Update(current.EndTime, current.Value)
            self.vwap_values.append(self.vwap.Current.Value)
            
            
    def VWAPReset(self):
        self.vwap.Reset()
    
    def BarHandler(self, consolidated):
        return