import pandas as pd
import numpy as np
from io import StringIO

class RedditStockSentiment(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020,8, 1)
        self.SetEndDate(2020, 9, 1)
        self.SetCash(100000)
        self.tickers = ["HD"]
        self.investP = 1/len(self.tickers)
        for stock in self.tickers:
            self.AddEquity(stock, Resolution.Hour)
        
        self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))
        self.SetBenchmark("HD")
        
        self.tradingday = 0
        
        csv = self.Download("https://www.dropbox.com/s/tb42cw10y2r4hd1/Foot_Traffic.csv?dl=1")
        self.df = pd.read_csv(StringIO(csv))
        self.df = self.df.drop([0, 1, 7, 8, 14, 15, 21, 22, 28, 29]) #Drops all weekend data
        
        self.Schedule.On(self.DateRules.EveryDay(), 
                 self.TimeRules.At(8, 30),        
                 self.runDaily) #Trade once a day at 8:30 am
        
        self.trafficValues = RollingWindow[int](5) #one trading week of foot traffic data
        self.smaSlope = RollingWindow[int](2) #two day rollingwindow for SMA

    def OnData(self, data):
        
        if self.trade != True: #Trade once a day at 8:30am
            return
        
        self.tradingday = (self.Time.day)-1
        
        for row in self.df.itertuples(): #Chooses data corresponding to today
            count = int(row[0])
            traffic = int(row[1])
            if count == self.tradingday:
                break
            else:
                continue
            
        self.trafficValues.Add(traffic) #Adds corresponding data to rolling window
    
        currentSMA = sum(self.trafficValues)/5 #Calculates the current SMA with a 5 day lookback
        
        self.smaSlope.Add(currentSMA) #Adds the SMA to a two day rolling window
        
        try:
            slope = (self.smaSlope[0] - self.smaSlope[1])/1 #Calculates the slope of the SMA  (SMA0-SMA1)/(DateToday-DateYesterday)
        except:
            slope = 0 #If error make slope 0. Only error if rolling window is missing data point
            
        self.Debug("Count: " + str(count) + " Day: " + str(self.tradingday) + " Slope: " + str(slope))
            
        if count == self.tradingday & count > 4 & (count != 7 or 14 or 21 or 28): #Double checking for current day, checking if atleast a week has passed for foot traffic data rolling window to fill, double checks to make sure we are not trading on weekend
            if(slope > 0) and not self.Portfolio["HD"].IsLong: #If slope is positive long
                self.SetHoldings("HD", self.investP, True)
            elif(slope < 0) and not self.Portfolio["HD"].IsShort: #If slope is negative short
                self.SetHoldings("HD", -(self.investP*0.4), True) #40% leverage
            elif(slope == 0): #If slope is 0 then liquidate. (Also known as the Flat insight)
                self.Liquidate("HD")
        
        self.trade = False
                
    def runDaily(self):
            self.trade = True