from AlgorithmImports import *

# QuantConnect algorithm to reproduce results in
# Zarattini, Carlo and Aziz, Andrew,
# Volume Weighted Average Price (VWAP) The Holy Grail for Day Trading Systems (November 13, 2023).
# Available at SSRN: https://ssrn.com/abstract=4631351 or http://dx.doi.org/10.2139/ssrn.4631351

class VWAPTrendTradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 11, 10)
        self.SetEndDate(2023, 11, 10)
        self.SetCash(100000)
        self.symbol = self.AddEquity("QQQ", Resolution.Minute).Symbol
        self.vwap = self.VWAP(self.symbol)

        self.Schedule.On(self.DateRules.EveryDay(self.symbol), 
                         self.TimeRules.BeforeMarketClose(self.symbol, 1), 
                         self.EndTrading)

        chart = Chart("Chart")
        chart.AddSeries(CandlestickSeries("QQQ", "$"))
        chart.AddSeries(Series("WVAP", SeriesType.LINE, "$", Color.Blue))
        chart.AddSeries(Series("Buys", SeriesType.SCATTER, "$", Color.Green, ScatterMarkerSymbol.TRIANGLE))
        chart.AddSeries(Series("Sells", SeriesType.SCATTER, "$", Color.Red, ScatterMarkerSymbol.TRIANGLE_DOWN))
        self.AddChart(chart)

    def EndTrading(self):
        self.ExitPosition()

    def OnData(self, data):
        if not data.ContainsKey(self.symbol) or data[self.symbol] is None:
            return

        self.plot("Chart", "Candles", data[self.symbol])
        price = data[self.symbol].Close

        if not self.vwap.IsReady:
            return

        self.Plot("Chart", "VWAP", self.vwap.Current.Value)

        if not self.Portfolio[self.symbol].Invested:
            self.EnterPosition(price)
        else:
            if self.Portfolio[self.symbol].IsLong and price < self.vwap.Current.Value:
                self.stopLoss()
                self.EnterPosition(price)
            elif self.Portfolio[self.symbol].IsShort and price > self.vwap.Current.Value:
                self.stopLoss()
                self.EnterPosition(price)

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            # Retrieve the fill price from the order event
            fillPrice = orderEvent.FillPrice
            
            # Check the order direction and plot accordingly
            if orderEvent.Direction == OrderDirection.Buy:
                self.Plot("Chart", "Buys", fillPrice)
            elif orderEvent.Direction == OrderDirection.Sell:
                self.Plot("Chart", "Sells", fillPrice)

    def EnterPosition(self, price):
        if self.vwap.Current.Value < price:
            self.SetHoldings(self.symbol, 1) # go long
        else:
            self.SetHoldings(self.symbol, -1) # go short

    def stopLoss(self):
        self.Liquidate(self.symbol)

    def ExitPosition(self):
        self.Liquidate(self.symbol)

    def OnEndOfDay(self, symbol: Symbol):
        pass