import numpy as np
from datetime import timedelta

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
class BasicTemplateAlgorithm(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2017,1, 1)  #Set Start Date
        self.SetEndDate(2018,4,1)    #Set End Date

        self.SetCash(100000)           #Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction)
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        self.num_stocks = 5
        self.big_number_of_stocks = 200
        
        #rolling window size for dollar volume calculation
        self.history_window = 200
        self.SetWarmUp( timedelta( self.history_window ) )
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 10), Action(self.rebalance))

        self.symbols = []
        self.volumes = {}
        self.prices = {}

    def CoarseSelectionFunction(self, coarse):
        '''reset symbol universe for price sorted by dollar volume desc'''
        selected = [x for x in coarse if (x.HasFundamentalData)]
        filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)[:self.big_number_of_stocks]
        syms = list( self.volumes.keys() )
        new_syms = []
        
        #add dollar volume and price info to our dict object
        for x in filtered:
            new_syms.append( x.Symbol )
            if x.Symbol not in syms:
                self.volumes[x.Symbol] = [x.DollarVolume]
                self.prices[x.Symbol] = [x.Price]
            else:
                self.volumes[x.Symbol].append( x.DollarVolume )
                self.prices[x.Symbol].append( x.Price )

            #keep the data the right length
            if len( self.volumes[ x.Symbol ] ) > self.history_window:
                self.volumes[ x.Symbol ] = self.volumes[ x.Symbol ][ :-1 ]
                self.prices[ x.Symbol ] = self.prices[ x.Symbol ][ :-1 ]
        
        #remove any not in current set
        syms = list( self.volumes.keys() )
        for sym in syms:
            if sym not in new_syms:
                del self.volumes[ sym ]
                del self.prices[ sym ]
        
        
        #make total dollar volume stat
        total_volumes = []
        syms_filtered = []
        for sym in new_syms:
            if len( self.volumes[ sym ] ) >= self.history_window: 
                total_volumes.append( np.sum( self.volumes[ sym ] ) )
                syms_filtered.append(sym)
         
        top_five_idx = np.argsort( total_volumes )[-self.num_stocks:]     
        
        self.symbols = [ syms_filtered[i] for i in top_five_idx]
        return self.symbols

    def OnData(self, data):
        pass
    
    def make_trades(self, w):
        for symbol in self.Portfolio.Keys:
            if ( symbol not in self.symbols ):
                self.SetHoldings(symbol, 0)
                
        self.Log(str(self.symbols))
        for i in range(0, self.num_stocks):
            self.AddEquity(self.symbols[i].Value)
            self.SetHoldings(self.symbols[i], w[i])
    
    def rebalance(self):
        if len(self.symbols) == self.num_stocks:
            if self.no_bear():
                return
            w = self.get_weights()
            self.make_trades(w)
    
    def no_bear( self ):
        spy_hist = self.History([self.spy], 120, Resolution.Daily).loc[str(self.spy)]['close']
        if self.Securities[self.spy].Price < spy_hist.mean():
            for symbol in self.Portfolio.Keys:
                if symbol.Value != "TLT":
                    self.Liquidate(symbol)
            self.AddEquity("TLT")
            self.SetHoldings("TLT", 1)
            return True
        else:
            return False
        
    def get_weights(self):
        return [1/self.num_stocks] * self.num_stocks

def take_top_stocks(self):
        hist = self.History(self.symbols, self.history_window, Resolution.Daily)
        
        self.price = {}
        for symbol in self.symbols:
            if str(symbol) in hist.index:
                this_vector = np.array( hist.loc[str(symbol)]['close'] )
                if len( this_vector ) == self.history_window:
                    self.price[symbol.Value] = np.array( this_vector )
            
        price_matrix = np.array([ self.price[symbol] for symbol in self.price.keys() ])
        
        lr = np.diff( np.log( price_matrix ) )
        # self.Log(str(lr))
        mu = np.mean( lr, axis=1)
        std = np.std( lr, axis=1)
        sharpe = mu / std
        # self.Log(str(sharpe))
        filt = list( np.nonzero( sharpe >= np.sort( sharpe )[ -self.tiny_num_stocks ] )[0] )
        syms = list( self.price.keys() )
        
        # return [ lr, syms ]
        return [ np.array( [ lr[i,:] for i in filt] ), [syms[i] for i in filt] ]