import numpy as np
from System.Drawing import Color
from decimal import Decimal

class movingAverageEnvPlotExample(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(2017, 6, 1)
        self.SetEndDate(2019, 1, 1)
        self.SetCash(5000)
        self.spy = self.AddEquity("SPY", Resolution.Daily)

        # Fixed deviation percentage for upper and lower band
        deviationPct = 0.05
        
        # Construct a simple moving average and a upper and lower band
        self.middleBand = self.SMA("SPY", 20, Resolution.Daily)
        self.upperBand = IndicatorExtensions.Times(self.middleBand,(1+deviationPct))
        self.lowerBand = IndicatorExtensions.Times(self.middleBand,(1-deviationPct))

        # Construct "Trade Plot"
        IndicatorPlot = Chart("Trade Plot")
        IndicatorPlot.AddSeries(Series("Upper Band", SeriesType.Line,"", Color.Blue))
        IndicatorPlot.AddSeries(Series("Middle Band", SeriesType.Line,"", Color.Green))
        IndicatorPlot.AddSeries(Series("Lower Band", SeriesType.Line,"", Color.Red))
        IndicatorPlot.AddSeries(Series("Close Price", SeriesType.Line,"", Color.Black))

        self.AddChart(IndicatorPlot)
        
    def OnData(self,data):
        if not self.upperBand.IsReady: return

        high = self.upperBand.Current.Value
        mid = self.middleBand.Current.Value
        low = self.lowerBand.Current.Value

        price = self.Securities["SPY"].Close    

        # Add values to "Trade Plot"
        self.Plot("Trade Plot", "Upper Band", str(high))
        self.Plot("Trade Plot", "Middle Band", str(mid))
        self.Plot("Trade Plot", "Lower Band", str(low))
        self.Plot("Trade Plot", "Close Price", str(price))

        # Enter a long position
        if not self.Portfolio.Invested:
                self.SetHoldings("SPY",1)