from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Data import *
from QuantConnect.Data.Market import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Securities import *
from QuantConnect.Orders import *
from datetime import datetime
from System.Drawing import Color

import decimal as d
import numpy as np

class Demo(QCAlgorithm):

    def Initialize(self):
        
        # configuration parameters (configurable inputs into the algorithm)
        MINUTES_AFTER_OPEN = 0
        MINUTES_BEFORE_CLOSE = 1
        SYMBOL = 'SPY'
        BBLENGTH = 20
        BBDEV = 2
        
        
        self.ORDER_MAP = ["Market", "Limit", "StopMarket", "StopLimit", "MarketOnOpen", "MarketOnClose", "OptionExercise"]
        

        # initialization
        self.SetStartDate(2020, 9, 1)
        self.SetEndDate(2020, 9, 3)
        self.SetCash(100000)
        self.stock = self.AddEquity(SYMBOL, Resolution.Minute)
        self.SetTimeZone("America/Chicago")
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.Schedule.On(self.DateRules.EveryDay(self.stock.Symbol), self.TimeRules.AfterMarketOpen(self.stock.Symbol, MINUTES_AFTER_OPEN), self.OnMarketOpen)
        self.Schedule.On(self.DateRules.EveryDay(self.stock.Symbol), self.TimeRules.BeforeMarketClose(self.stock.Symbol, MINUTES_BEFORE_CLOSE), self.OnMarketClose)
        # set the trade flag to False. we'll only start trading when the flag flips to True (after the market open event)
        self.tradeFlag = False
        
        #self.bb = self.BB(SYMBOL, BBLENGTH, 2, MovingAverageType.Exponential, Resolution.Minute)
        self.bb = BollingerBands(BBLENGTH, 2, MovingAverageType.Exponential)


        self.lastPrice = 0.0
        self.pl = 0.0
        self._lo = None
        
        self.SetWarmUp(BBLENGTH * 5, Resolution.Minute)     ## Change set warm up
        
        self.Consolidate(SYMBOL, timedelta(minutes=5), self.OnStockBarConsolidated)
        
        

    def OnMarketOpen(self):
        self.tradeFlag = True
        
    def OnMarketClose(self):
        # liquidate all holdings
        if self.stock.Invested:
            self.Liquidate(self.stock.Symbol, "EOD Liquidate")
        else:
            self.Transactions.CancelOpenOrders()

        # reset trade flag for following day
        self.tradeFlag = False


    def OnStockBarConsolidated(self, consolidated):
        #if self.IsWarmingUp or not self.tradeFlag or not self.bb.IsReady:      ######### Need to remove
        #    return
        
        self.bb.Update(consolidated.EndTime, consolidated.Close)

        if not self.bb.IsReady:
            return

        self.Plot("BB", "Upper", self.bb.UpperBand.Current.Value)
        self.Plot("BB", "Price", consolidated.Close)
        self.Plot("BB", "Lower", self.bb.LowerBand.Current.Value)