from QuantConnect.Data.UniverseSelection import *
import pandas as pd
import numpy as np

class EnhancedShortTermMeanReversionAlgorithm(QCAlgorithm):

    def Initialize(self):

        # Note - self is the object, it has methods (functions) and properties (attributes/variables)

        # We define start date, end date and cash by using methods belonging to the object        
        self.SetStartDate(2005,1,1) #Set Start Date
        self.SetEndDate(2018,1,27)  #Set End Date       
        self.SetCash(50000)         #Set Strategy Cash

        # Need to use the universe selection model to select assets for our Strategy
        self.UniverseSettings.Resolution = Resolution.Daily # self.UniverseSettings.Resolution property is the resolution of all the assets added to Universe
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction) # use the AddUniverse() method to add a universe
        # takes function filters as parameters, hence we define the functions below
        # - coarse selection function: accepts an IEnumerable<CoarseFundamental> and returns an IEnumerable<Symbol>
        # - fine selection function: accepts an IEnumerable<FineFundamental> and returns an IEnumerable<Symbol>
        
        # benchmark
        self.AddEquity("SPY", Resolution.Minute) # Add minute resolution data for SPY benchmark using AddEquity() method 
        
        # Rebalancing
        # rebalance the universe selection once a month
        self.rebalence_flag = 0
        # make sure to run the universe selection at the start of the algorithm even it's not the month start
        self.first_month_trade_flag = 1
        self.trade_flag = 0  
        
        # Number of quantiles for sorting returns for mean reversion
        self.nq = 5 # setup a property of self called nq and set it to 5
        # Number of quantiles for sorting volatility over five-day mean reversion period
        self.nq_vol = 3 # setup a property of self called nq_vol and set it to 3
        # the symbol list after the coarse and fine universe selection
        self.universe = None # Initialize the universe to None

        # Scheduling        
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 303), Action(self.get_prices))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 302), Action(self.daily_rebalance))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 301), Action(self.short))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 300), Action(self.long))
    
    def monthly_rebalance(self):
        # rebalance the universe every month
        self.rebalence_flag = 1
 
    def CoarseSelectionFunction(self, coarse): # coarse is the only attribute other than self, I assume it is the list to coarse filter
        
        if self.rebalence_flag or self.first_month_trade_flag: # if we have to rebalance or it is the very first month (needed to run universe selection at start of algorithm)
            # drop stocks which have no fundamental data or have too low prices
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)] # list comprehension
            # if x object (stock) has property HasFundamentalData = True and price greater than 5, gives a new list
            # rank the stocks by dollar volume and choose the top 50
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) # use lambda function to reverse sort the selected list by DollarVolume property

            return [ x.Symbol for x in filtered[:50]] # give the 50 biggest symbols by DollarVolume
        else:
            return self.universe # otherwise return the previous universe

    def FineSelectionFunction(self, fine): # I assume FineSelectionFunction will filter the results of the CoarseSelectionFunction

        if self.rebalence_flag or self.first_month_trade_flag:
            # filter the stocks which have positive EV To EBITDA (similar to PE ratio)
            filtered_fine = [x for x in fine if x.ValuationRatios.EVToEBITDA > 0] # use list comprehension to only get stocks with positive EVToEBITDA
            self.universe = [x.Symbol for x in filtered_fine] # to get the universe just take the symbol property of each stock
            
            self.rebalence_flag = 0 # set the rebalence flag to off
            self.first_month_trade_flag = 0 # set first month flag to off
            self.trade_flag = 1 # set trade flag to on
            
        return self.universe # return the universe after coarse and fine filtering
        

    def OnData(self, data):
        pass
    
    def short(self):
        if self.universe is None: return
        SPY_Velocity = 0
        self.long_leverage = 0
        self.short_leverage = 0
        # request the history of benchmark
        pri = self.History(["SPY"], 200, Resolution.Daily)
        pos_one = (pri.loc["SPY"]['close'][-1])
        pos_six = (pri.loc["SPY"]['close'][-75:].mean())
        # calculate velocity of the benchmark 
        velocity_stop = (pos_one - pos_six)/100.0
        SPY_Velocity = velocity_stop
        if SPY_Velocity > 0.0:
            self.long_leverage = 1.8
            self.short_leverage = -0.0
        else:
            self.long_leverage = 1.1
            self.short_leverage = -0.7
        for symbol in self.shorts:
            if len(self.shorts) + self.existing_shorts == 0: return
            self.AddEquity(symbol, Resolution.Minute)
            self.SetHoldings(symbol, self.short_leverage/(len(self.shorts) + self.existing_shorts))                                
 
    def long(self):
        if self.universe is None: return
        for symbol in self.longs:
            if len(self.longs) + self.existing_longs == 0: return
            self.AddEquity(symbol, Resolution.Minute)
            self.SetHoldings(symbol, self.long_leverage/(len(self.longs) + self.existing_longs))                                
       
    def get_prices(self):
        if self.universe is None: return
        # Get the last 6 days of prices for every stock in our universe
        prices = {}
        hist = self.History(self.universe, 6, Resolution.Daily)
        for i in self.universe:
            if str(i) in hist.index.levels[0]:
                prices[i.Value] = hist.loc[str(i)]['close']
        df_prices = pd.DataFrame(prices, columns = prices.keys()) 
        # calculate the daily log return
        daily_rets = np.log(df_prices/df_prices.shift(1))
        # calculate the latest return but skip the most recent price
        rets = (df_prices.iloc[-2] - df_prices.iloc[0]) / df_prices.iloc[0]
        # standard deviation of the daily return
        stdevs = daily_rets.std(axis = 0)
        self.ret_qt = pd.qcut(rets, 5, labels=False) + 1
        self.stdev_qt = pd.qcut(stdevs, 3, labels=False) + 1
        self.longs = list((self.ret_qt[self.ret_qt == 1].index) & (self.stdev_qt[self.stdev_qt < 3].index))
        self.shorts = list((self.ret_qt[self.ret_qt == self.nq].index) & (self.stdev_qt[self.stdev_qt < 3].index))

 
    def daily_rebalance(self):
        # rebalance the position in portfolio every day           
        if self.universe is None: return
        self.existing_longs = 0
        self.existing_shorts = 0
       
        for symbol in self.Portfolio.Keys:
            if (symbol.Value != 'SPY') and (symbol.Value in self.ret_qt.index):
                current_quantile = self.ret_qt.loc[symbol.Value]
                if self.Portfolio[symbol].Quantity > 0:
                    if (current_quantile == 1) and (symbol not in self.longs):
                        self.existing_longs += 1
                    elif  (current_quantile > 1) and (symbol not in self.shorts): 
                        self.SetHoldings(symbol, 0)
                elif self.Portfolio[symbol].Quantity < 0:
                    if (current_quantile == self.nq) and (symbol not in self.shorts):
                        self.existing_shorts += 1
                    elif (current_quantile < self.nq) and (symbol not in self.longs): 
                        self.SetHoldings(symbol, 0)