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(2013,1,1)  #Set Start Date
        self.SetEndDate(2013,2,1)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        self.AddEquity("SPY", Resolution.Daily)
        
        # what resolution should the data *added* to the universe be?
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # How many stocks in the starting universe?
        self.__numberOfSymbols = 20
        
        # How many stocks in the portfolio?
        self.number_stocks = 5
        
        # this add universe method accepts two parameters:
        self.AddUniverse(self.CoarseSelectionFunction)
        
        # How far back are we looking for momentum?
        self.momentum_period = 20
        
        # Schedule Indicator Update, Ranking + Rebal
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), Action(self.rebalance))
        
        # Set Risk Factor for position sizing
        self.risk_factor = 0.001
        
        # Set empty list for universe
        self.universe = []
        
        # Set empty dictionary for managing & ranking the slope
        self.indicators = {}
        
        # Set empty dictionary to manage ATR position sizing for ranked securities
        self.target_portfolio = {}

    def OnData(self, data):
        pass
    
    # 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 today.day == 1 and (today.month == 1 or today.month == 4 or today.month == 7 or today.month == 10):
            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 rebalance(self):
        
        # This updates the indicators at each data step(based on resolution)
        for symbol in self.universe:
            
            # is symbol iin Slice object? (do we even have data on this step for this asset)
            if not self.Securities.ContainsKey(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
            
            # Update the dictionary for the indicator
            #history = self.History(symbol, self.momentum_period, Resolution.Daily)["close"].unstack(level=0)
            #self.indicators[symbol] = regression_slope(history)
            self.indicators[symbol] = update_indicator(self.Time, self.Securities[symbol].Price)
            
        # Sort the Dictionary from highest to lowest and take the top values
        self.target_portfolio = sorted(self.indicators, key=self.indicators.get, reverse=True)[:self.number_stocks]
            
        for symbol in self.target_portfolio:
            
            # Liquidate if not in target portfolio
            if symbol.Invested:
                if symbol not in self.target_portfolio:
                    self.Liquidate(symbol)
                elif symbol in self.target_portfolio:
                    continue
                
            # Place orders for stocks not current in target portfolio
            elif symbol.Invested == False:
                
                #self.Log("{} {} {}"symbol,)
                
                # Update ATR for the stock in the new dictionary
                self.target_portfolio[symbol].update_value(self.Time, self.Securities[symbol].Price)
                
                self.Log("{} {} {}".format(symbol,self.Securities[symbol].Price,self.indicators[symbol].ATR.Current.Value))
                
                # Calculate number of shares to buy
                quantity = self.target_portfolio[symbol].ATR.Current.Value/(Portfolio.TotalPortfolioValue*self.risk_factor)
                
                # Send Orders
                self.MarketOrder(stock, quantity)
            
            
class SymbolData(object):
    def __init__(self, symbol, context):
        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.ATR = context.ATR(symbol, 20)

    """
    Runtime Error: Python.Runtime.PythonException: NotSupportedException : AverageTrueRange does not support Update(DateTime, decimal) method overload. Use Update(IBaseDataBar) instead.
    """
        
        # Use the below update method for ATR which is used for position sizing
    def update_value(self, time, value):
        self.ATR.Update(time, value)
        
class CustomIndicator(object):
    def __init__(self, symbol, context):
        history = self.History(symbol, self.momentum_period, Resolution.Daily)["close"].unstack(level=0)

        # Get Regression Stats & Compute Slope * R^2
    def regression_slope(self, time, value):  
        x = np.arange(len(history))
        log_ts = np.log(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)
        # Return the Indicator
        return annualized_slope
        
    def update_indicator(self, time, value):
        self.regression_slope.Update(time, value)