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

class RedditStockSentiment(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017,1, 10)  # Set Start Date
        self.SetEndDate(2020, 9, 14) #Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.tickers = ["SPCE", "LULU", "CCL", "SDC", "NKLA", "TSLA", "F", "BA", "PTON", "CHWY", "RKT"] #Stocks picked by briefly looking through the wallstreetbets reddit and finding frequently mentioned stocks
        self.investP = 1/len(self.tickers) #Equal weight portfolio
        self.SetTimeZone(TimeZones.Chicago) #Set to Chicago time
        
        for stock in self.tickers:
            self.AddEquity(stock, Resolution.Hour) #Sets resolution to hour bars
        
        self.AddRiskManagement(TrailingStopRiskManagementModel(0.08)) #Risk management
        
        self.trade = True #OnData will run when the program when the program is first executed
        
        csv = self.Download("https://www.dropbox.com/s/ibd2zty0ytoqsos/The_Wall_Street_Bets_Sentiment.csv?dl=1") #Downloads data
        self.df = pd.read_csv(StringIO(csv)) #Read into a dataframe
        
        self.Schedule.On(self.DateRules.EveryDay(), 
                 self.TimeRules.At(8, 30),        
                 self.runDaily) #Runs runDaily (sets self.trade to True) at 8:30am Chicago time
        
    def OnData(self, data):
        
        if self.trade != True: #Runs daily at 8:30am Chicago time
            return
        
        #Assigns three variables with the current time(year, month, day)
        algYear = self.Time.year
        algMonth = self.Time.month
        algDay = self.Time.day
        
        #Iterates through the dataframe tuples
        for row in self.df.itertuples():
            try: #If there is an error with dates a tuple(row) then continue to the next row
                date = row[-1] #Date is the index of the tuple(final colum in the dataframe)
                
                #Date is parsed as a string and assigned to three variables
                year = date[0:4]
                month = date[5:7]
                day = date[8:10]
            except:
                continue
            
            #Compares current date with the date from the tuple
            if (int(year) != algYear) or (int(month) != algMonth) or (int(day) != algDay):
                continue
            
            '''
            averageSentiment is calculated by averaging the sentiment of the submission title, body text, and all comments.
            This process can be very time consuming. 
            '''
            
            #Assigns several variables with data from the tuple
            stock = str(row[1])
            averageSentiment = float(row[4]) 
            numberOfComments = int(row[3])
            score = int(row[7])
            
            adjust = 1
            
            #If score is low, allocate a smaller percentage to the position taken for the underlying equity
            if score <= 20:
                adjust = 0.5
        
            if(averageSentiment >= 0.1) and not self.Portfolio[stock].IsLong:
                self.SetHoldings(stock, self.investP*adjust, True)
            if(averageSentiment <= -0.1) and not self.Portfolio[stock].IsShort:
                self.SetHoldings(stock, -(self.investP*0.4)*adjust, True) #40% leverage
            if (-0.1 < averageSentiment < 0.1):
                self.Liquidate(stock)
                
            self.trade = False
                
    def runDaily(self):
        self.trade = True