import numpy as np
import datetime
from scipy import stats 

### <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 StocksOnTheMove(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(2016,1,1)  #Set Start Date
        self.SetEndDate(2016,1,30)    #Set End Date
        self.SetCash(300000)           #Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        self.AddEquity("SPY", Resolution.Minute)
        
        # what resolution should the data *added* to the universe be?
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        #self.UniverseSettings.Resolution = Resolution.Daily        
        self.UniverseSettings.Leverage   = 1
        
        # How many stocks in the starting universe?
        #self.__numberOfSymbols = 700
        self.__numberOfSymbols = 500
        
        # How many stocks in the portfolio?
        self.number_stocks = 21
        
        # this add universe method accepts two parameters:
        self.AddUniverse(self.CoarseSelectionFunction)
        
        # How far back are we looking for momentum?
        self.momentum_period = 66
        
        # Set ATR window
        self.atr_window = 10
        self.SetWarmUp(self.atr_window)
        
        # Schedule Indicator Update, Ranking + Rebal
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), 
        #                 self.TimeRules.AfterMarketOpen("SPY", 30), 
        #                 Action(self.rebalance))
                         
        self.Schedule.On(self.DateRules.EveryDay("SPY"), 
                         self.TimeRules.BeforeMarketClose("SPY", 0), 
                         Action(self.UpdateIndicators))
        
        
        # Set empty list for universe
        self.universe = []
        
        # Set empty dictionary for managing & ranking the slope
        self.indicators_r2 = {}
        self.indicators = {}
        self.atr = {}
        
        self.splotName = 'Strategy Info'
        sPlot = Chart(self.splotName)
        sPlot.AddSeries(Series('Leverage',  SeriesType.Line, 0))
        self.AddChart(sPlot)
        
        self.last_month_fired_coarse    = None #we cannot rely on Day==1 like before
        self.last_month_fired_rebalance = None #we cannot rely on Day==1 like before
        
    def UpdateIndicators(self):
        
        #self.Log("UpdateIndicators: Started Function")
        # This updates the indicators at each data step
        for symbol in self.universe:
            
            # is symbol iin Slice object? (do we even have data on this step for this asset)
            if self.Securities.ContainsKey(symbol):
                # Update the dictionary for the indicator
                if symbol in self.indicators_r2:
                    self.indicators_r2[symbol].update(self.Securities[symbol].Price)
                    
        self.account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot(self.splotName,'Leverage', float(self.account_leverage))
        self.Log("UpdateIndicators: Finished Successfully")
        
    # Run a coarse selection filter for starting universe
    def CoarseSelectionFunction(self, coarse):
    
        today = self.Time
        #self.Log("Day = {} Month = {}".format(today.day,today.month))
        
        # Set the Universe to rebalance on the 1st day of each quarter (can play around with this as required)
        if self.last_month_fired_coarse != today.month and (today.month == 1 or today.month == 4 or today.month == 7 or today.month == 10):
            self.last_month_fired_coarse = today.month
            
            self.Log("Day = {} Month = {}".format(today.day,today.month))
            
            CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
            sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True)
            result = [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]
            self.universe = result
            return self.universe
        else:
            return self.universe
            
    def OnSecuritiesChanged(self, changes):
        
        self.Log("OnSecuritiesChanged: Starting Removing Securities that left universe")
        # Delete indicator from the dict to save Ram
        for security in changes.RemovedSecurities:
            if security.Symbol in self.indicators_r2:
                del self.indicators_r2[security.Symbol]
                self.Liquidate(security.Symbol)
        
        self.Log("OnSecuritiesChanged: Removed Securities that have left universe")        
        # Init a new custom indicator
        for security in changes.AddedSecurities:
            self.indicators_r2[security.Symbol] = RegressionSlope(self, security.Symbol, self.momentum_period,  Resolution.Daily)

        self.Log("OnSecuritiesChanged: Added Indicator for new securities added to universe")
        self.Log("OnSecuritiesChanged:  Finished Successfully")
                
    def OnData(self, data):
        
        #self.target_portfolio = None
        today = self.Time
        if self.last_month_fired_rebalance != self.last_month_fired_coarse:
            # ensure we are fireing after coarse
            self.last_month_fired_rebalance = self.last_month_fired_coarse
            
            self.Log("OnData: Start Rebalance")
        
            # get values from dict
            symbols, slopes = zip(*[(symbol, self.indicators_r2[symbol].value) \
                                       for symbol in self.indicators_r2 \
                                           if self.indicators_r2[symbol].value is not None])
            
            # sort 
            idx_sorted = np.argsort(slopes)[::-1] # [::-1] slices backwards i.e. flips to reverse the sort order
            symbols = np.array(symbols)[idx_sorted]
            slopes = np.array(slopes)[idx_sorted]
            
            #self.Log("Finished Generating Slopes")
            
            # Sort the Dictionary from highest to lowest and take the top values
            self.target_portfolio = []
            self.target_portfolio = symbols[:self.number_stocks]
            #self.atr = symbols
            #self.Log(str(self.target_portfolio))
            self.Log("LENGTH TARGET PORTFOLIO: {}:".format(str(len(self.target_portfolio))))
            self.Log("TYPE TARGET PORTFOLIO: {}:".format(str(type(self.target_portfolio))))
            #self.Log("TARGET PORTFOLIO: {}:".format(str(self.target_portfolio)))
            
            # Get ATR for the symbol
            for symbol in self.target_portfolio:
                #self.Log("Success looping through new target portfolio")
                # is symbol in Slice object? (do we even have data on this step for this asset)
                if not data.ContainsKey(symbol):
                    #del self.target_portfolio[symbol]
                    self.Log("Symbol with No Data: {}:".format(str(self.target_portfolio[symbol])))
                    continue
                
                
                # 686 | 13:35:43: Runtime Error: Python.Runtime.PythonException: AttributeError : 'NoneType' object has no attribute 'Price'
                if data[symbol] is None:
                    continue
                
                # Does this slice have the price data we need at this moment?
                if data[symbol].Price is None:
                    continue
                
                #self.Log("OnData: Finished Data Quality Checks")
                
                #i = 0
                #self.Log("SYMBOL: {}:".format(str(symbol)))
                #self.atr[symbol] = SymbolData(symbol, self, self.atr_window).get_atr()
                
                if symbol not in self.indicators:
                    self.indicators[symbol] = SymbolData(symbol, self, self.atr_window)
                    
                #self.atr[symbol] = self.indicators[symbol].get_atr()
                self.indicators[symbol].update(data[symbol])
                
                #self.Log("OnData: Finished Updating ATRs for Target Portfolio")
                
                if self.IsWarmingUp: continue
            
                #self.atr[symbol] = float(self.indicators[symbol].get_atr()/self.Securities[symbol].Price)
                self.atr[symbol] = self.indicators[symbol].get_atr()
                #i += 1
                #self.atr[symbol] = self.target_portfolio.get_atr()
                #self.atr[i] = SymbolData.get_atr(i, self)
                #self.atr[i].get_atr(i, self)
                #continue
                
            #self.Log("SELF ATR: {}:".format(str(self.atr)))
            
            '''
            Seems that self.target_portfolio is generating QC symbols rather than a list of tickers as string & hence is not indexable. Not sure how to fix yet.
            I'm also not sure if ATR would need updating as I only need it once per month for rebalance
            '''
            
            #self.Log("TARGET PORTFOLIO: {} : ATRS : {}".format(str(self.target_portfolio,self.atr[symbol])))
            
            # new symbol? setup indicator object. Then update
            #if symbol not in self.indicators:
            #    self.indicators[symbol] = SymbolData(symbol, self, self.atr_window)
            #self.indicators[symbol].update(data[symbol])
            
            # Enter or exit positions
            self.Log("OnData: Begin Enter/Exit of Positions:")
            
            for symbol in self.universe:
            #for symbol in self.Portfolio.
            
            #for symbol in self.target_portfolio:
                
                
                #k = 1/np.sum(self.atr)
                # Case: invested in the current symbol
                if self.Portfolio[symbol].HoldStock:
                    # Exit if not a target aset
                    if symbol not in self.target_portfolio:
                        self.Liquidate(symbol)
                        self.Log("OnData: Trying to remove security from indicators list")
                        del self.indicators[symbol]
                        self.Log("OnData: Success Removing")
                    elif symbol in self.target_portfolio:
                        # Add Rebalance
                        k = float(1/sum(self.atr.values()))
                        self.SetHoldings(symbol, k/float(self.atr[symbol]))
                        continue
                        
                    
                # Case: not invested in the current symbol
                else:
                    '''
                    k = float(1/sum(self.atr.values()))
                    self.Log("K VALUE: {}".format(k))
                    self.SetHoldings(symbol, k/float(self.atr[symbol]))
                    '''
                    
                    # symbol is a target, enter position
                    if symbol in self.target_portfolio:
                        
                        # Update ATR for the stock in the new dictionary
                        #self.atr[symbol].update_bar(self.Time, data[symbol].Price)
                        
                        #self.Log("SYM: {}, PRICE: {}, ATR: {}, R2: {}".format(symbol, self.Securities[symbol].Price, float(self.atr[symbol]), self.indicators_r2[symbol].value))
                        #self.Log("{}, {}, {}, {}".format(symbol, self.Securities[symbol].Price, float(self.atr[symbol]), self.indicators[symbol].value, self.indicators_r2[symbol].value))                        
                        # Send Orders - Equal Weighted
                        #self.SetHoldings(symbol, 0.99/float(self.number_stocks))
                        # Send Orders - Volatility Weighted
                        k = float(1/sum(self.atr.values()))
                        #self.Log("K VALUE: {}".format(k))
                        self.SetHoldings(symbol, k/float(self.atr[symbol]))
                    

class RegressionSlope():
    
    def __init__(self, algo, symbol, window, resolution):
        # set up params of per-asset rolling metric calculation
        self.symbol = symbol
        self.window = window
        self.resolution = resolution
        
        # the value we access, None until properly calulated
        self.value = None
        
        # We will store the historical window here, and keep it a fixed length in update
        self.history = []
        
        # download the window. Prob not great to drag algo scope in here. Could get outside and pass in.
        hist_df = algo.History([symbol], window, self.resolution)
        
        # Case where no data to return for this asset. New asset?
        if 'close' not in hist_df.columns:
            return
        
        # store the target time series
        self.history = hist_df.close.values
        
        # calulate the metrics for the current window
        self.compute()

    def update(self, value):
        # update history, retain length
        self.history = np.append(self.history, float(value))[1:]
        
        # calulate the metrics for the current window
        self.compute()
    
    def compute(self):
        
        # Case where History faiiled to return window, waiting to acrew
        # prevent calc until window is statisfied
        if len(self.history) < self.window:
            return
        
        # copied from previous
        x = np.arange(len(self.history))
        log_ts = np.log(self.history)
        slope, intercept, r_value, p_value, std_err = stats.linregress(x, log_ts)
        annualized_slope = (np.power(np.exp(slope), 250) - 1) * 100
        annualized_slope = annualized_slope * (r_value ** 2)
        
        # update value
        self.value = annualized_slope
        
class SymbolData(object):
    def __init__(self, symbol, context, window):
        self.symbol = symbol
        """
        I had to pass ATR from outside object to get it to work, could pass context and use any indica
        var atr = ATR(Symbol symbol, int period, MovingAverageType type = null, Resolution resolution = null, Func`2[Data.IBaseData,Data.Market.IBaseDataBar] selector = null)
        """
        self.window    = window
        self.indicator = context.ATR(symbol, self.window)
        self.atr       = 0.0

    """
    Runtime Error: Python.Runtime.PythonException: NotSupportedException : AverageTrueRange does not support Update(DateTime, decimal) method overload. Use Update(IBaseDataBar) instead.
    """
    def update(self, bar):
        self.indicator.Update(bar)
            
    def get_atr(self):
        return self.indicator.Current.Value