from QuantConnect.Python import PythonQuandl

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2003,1, 1)  #Set Start Date
        self.SetEndDate(2018,8,5)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        self.AddEquity("SPY", Resolution.Daily)
        self.AddEquity("AGG", Resolution.Daily)
        self.qGDP = "FRED/GDP"
        self.qNCBCEL = "FRED/NCBCEL"
        Quandl.SetAuthCode("YjQrjz6zF8mVWdGTz3fa")
        self.AddData(Fred, self.qGDP, Resolution.Daily, TimeZones.NewYork, True)
        self.AddData(Fred, self.qNCBCEL, Resolution.Daily, TimeZones.NewYork, True)
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(23, 0), self.rebalance)
        self.months = -1
        self.quarterly_rebalance = False
        #to help notice trends
        self.prev_value = 0

    def rebalance(self):
        #quarterly rebalance
        self.months+=1
        if self.months%3 == 0:
            self.quarterly_rebalance = True
        if self.quarterly_rebalance:
            indicator = float(self.Securities[self.qNCBCEL].Price)/float(self.Securities[self.qGDP].Price) 
            self.Plot("Buffet Indicator","Indicator", indicator)
            if indicator <= 1:
                if self.prev_value < indicator:
                    if self.Portfolio["AGG"].Invested:
                        self.Liquidate()
                    self.SetHoldings("SPY",1)
            else:
                if self.Portfolio["SPY"].Invested:
                        self.Liquidate()
                self.SetHoldings("AGG",1)
            self.quarterly_rebalance = False
            self.prev_value = indicator
    
    def OnData(self, data):
        pass

class Fred(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = 'value'