import numpy as np
from collections import deque

class RSharpeStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 12, 31) # Set End Date
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.UniverseSettings.Leverage = 1
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFilter)
        
        self.coarse_count = 10
        self.max_holding = timedelta(days=50) # max holding day
        self.delay_reopen = timedelta(days=10) # reopen after X days
        self.RS = { }; # buffer of rolling sharpe indicator for each symbol
        self.traded = { };# register dates of traded securites
        #parameters for rolling sharpe ratio
        self.sharpe_period=90
        self.RoC_max=0.5
        self.sharpe_min_avg=0.1
        self.sharpe_min_hold=0.5
        self.rank_top=5 # the top X of the rank
        self.num_stocks_min=5000 # minimum number of ready stocks in all stocks
        
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        self.SetWarmup(self.sharpe_period)
    
    # Filter function    
    def CoarseSelectionFilter(self, coarse):
        avg_sharpe = 0 # average rolling sharpe of all stocks
        num_stocks = 0 # number of ready stocks
        for stk in coarse:
            if stk.Symbol not in self.RS :
                self.RS[stk.Symbol] = RollingSharpe(symbol=stk.Symbol,period=self.sharpe_period)
                ################################################################################
                # history warmup may hold one iteration for more than 10 mins causing error stop
                if stk.Price > 2 and False:
                    # initialize the indicator with the daily history close price
                    history = self.History([stk.Symbol], timedelta(days=self.sharpe_period), Resolution.Daily)
                    if len(history)>0 :
                        for time, row in history.loc[stk.Symbol].iterrows():
                            self.RS[stk.Symbol].Update(time, row["close"])
            #else:
            ######
                # Updates indicators
            self.RS[stk.Symbol].Update(time=stk.EndTime, price=stk.Price)
            if stk.Price > 2 and self.RS[stk.Symbol].IsReady :
                avg_sharpe+=self.RS[stk.Symbol].Value
                num_stocks+=1
        if avg_sharpe!=0 and num_stocks!=0:        
            avg_sharpe/=num_stocks
        
        # Filter the values of the dict:
        # 1. bull's market: average rolling Sharpe ratio of all stocks > sharpe_min
        # 2. price > 2
        # 3. last sharpe_period's performance < RoC_max
        if num_stocks >= self.num_stocks_min and avg_sharpe > self.sharpe_min_avg :
            new = list(filter(lambda x: x.price>2 and x.IsReady and x.roc.Current.Value<self.RoC_max, self.RS.values()))
            # Sorts the values of the dict in descending order
            new.sort(key=lambda x: x.Value, reverse=True)
        else: # it's bear's market
            new=[ ]
        
        # check current holdings to close
        old=[ ]
        if self.Portfolio.Invested :
            for stk in self.Portfolio.Values:
                if self.Portfolio[stk.Symbol].Price > 2 and self.Portfolio[stk.Symbol].Invested and\
                    self.Time - self.RS[stk.Symbol].LastOpen < self.max_holding and\
                    self.RS[stk.Symbol].Value >= self.sharpe_min_hold:
                    old.append(self.RS[stk.Symbol])
                else:
                    self.RS[stk.Symbol].LastClose = self.Time
                
        # securites number control
        if len(new)>0:
            if len(old)>0:
                new = list(filter(lambda x: x not in old, new))
            new = new[:min(self.rank_top,self.coarse_count-len(old))]
            for x in new:
                if self.Time - x.LastClose >= self.delay_reopen:
                    x.LastOpen = self.Time
                    self.Log('New symbol: ' + str(x.Symbol) + ' Sharpe: current ' + str(round(x.Value,2)) + ' AvgRoll ' + str(round(avg_sharpe,2)))
        
        self.Log('AvgSharpe: '+str(round(avg_sharpe,2))+' Ready Stocks: '+str(num_stocks)+ ' total '+str(len(self.RS))+' len: new '+str(len(new))+' old '+str(len(old)))
        return [ x.Symbol for x in old+new ]
        
    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)

        # set equal allocation in each security in our universe
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 1/self.coarse_count/2)

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)

class RollingSharpe(object):
    def __init__(self, symbol, period):
        self.Symbol = symbol
        self.Time = datetime.min
        self.Value = 0
        self.price = 0
        self.IsReady = False
        self.period = period
        self.input = deque(maxlen=period) # limite number of inputs
        self.roc = RateOfChange(symbol, period)
        self.LastOpen = datetime(2000, 1, 1)
        self.LastClose = datetime(2000, 1, 1)
        

    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)
    
    # calculate Sharpe ratio for series x
    def sharpe(self, x, r = 0, scale = np.sqrt(250)):
        x=np.asfarray(x)
        if  x.ndim > 1 : # more than 1 dimention (1 column)
            print("x is not a vector or univariate time series")
            return
        if np.isnan(x).any() :
            print("NaNs in x")
            return
        if x.shape[0] == 1: # only 1 row
            return(np.nan)
        else :
            if len(x)<self.period:
                return(np.nan)
            else:
                y = np.diff(x)
        return(scale * (np.mean(y) - r)/np.std(y))
            
    def Update(self, time, price):
        self.price = price
        # append to the right because np.diff() uses right to minus left
        self.input.append(price)
        self.Time = time
        # calculate last x period's Sharpe ratio
        self.Value = self.sharpe(self.input)
        self.roc.Update(time, price)
        self.IsReady = len(self.input) == self.input.maxlen and self.roc.IsReady