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(2012,1,1)  #Set Start Date
        self.SetEndDate(2012,2,5)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        self.atrWindow = 20
        self.lookBackWindow = self.atrWindow + 5
        self.riskFactor = 0.001
        self.momentumWindowLength = 90
        self.rankTablePercentile = .2
        self.significantPositionDifference = 0.1
        self.numberOfSymbols = 20
   
        self.atrDict = {}
        self.AddEquity('SPY', Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
       
      
        self.Schedule.On(self.DateRules.MonthStart("SPY"), 
                         self.TimeRules.AfterMarketOpen("SPY", 60), 
                         Action(self.Rebalance))
                         
    
    def CoarseSelectionFunction(self, coarse):
        CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
        sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True)
        selectedSecurites = [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbols]]
        return selectedSecurites
                
    def slope(self,ts):
        x = np.arange(len(ts))
        slope, intercept, rValue, pValue, std_err = stats.linregress(x, ts)
        annualizedSlope = (1 + slope)**250
        return annualizedSlope * (rValue ** 2)    
        
    def Rebalance(self):
        numberOfSecurities = len(self.Securities.Keys)
        self.Debug("Number of securities in universe: "+str(numberOfSecurities))
        if numberOfSecurities <= 1: return
        history = self.History(self.Securities.Keys,  self.lookBackWindow, Resolution.Daily)
        for symbol in history.index.get_level_values(0):
            self.atrDict[symbol]=self.ATR(symbol,self.atrWindow)
       
        closes = history["close"].unstack(level=0)
        #Calculate slopes for each asset
        slopes = np.log(closes.tail(self.momentumWindowLength)).apply(self.slope)
        rankingTable = slopes[slopes > slopes.quantile(1 - self.rankTablePercentile)].sort_values(ascending=False)
        self.Debug("Top 5 in ranking table: ")
        self.Debug(rankingTable.head())
        estimatedCashBalance = float(self.Portfolio.Cash)
        # close positions that are no longer in the top of the ranking table
        positions =  self.Portfolio.Values
        for security in positions:
            if security not in rankingTable.index:
                self.Liquidate(security.Symbol)
            else:
                positionSize = positions[security].Amount
                newPositionSize = self.GetPositionSize(security)
                if position_size<= 0:
                    break
                if self.SignificantChangeInPositionSize(newPositionSize, positionSize):
                    estimatedCost = price * (newPositionSize - positionSize)
                    self.MarketOrder(security, newPositionSize, self.Securities[security].Price)
                    estimatedCashBalance -= estimatedCost
        #Open new positions based on ranking_table
        for security in rankingTable.index:
            if security not in positions:
                newPositionSize = self.GetPositionSize(security)
                if newPositionSize <= 0:
                    continue
                estimated_cost = float(self.Securities[security].Price) * newPositionSize
                if estimatedCashBalance > estimatedCost:
                    self.MarketOrder(security, newPositionSize)
                    estimatedCashBalance -= estimatedCost
        
    def GetPositionSize(self,symbol):
        averageTrueRange = self.atrDict[symbol].Current.Value
        self.Debug("Average true range: "+str(averageTrueRange))
        positionSize = 0
        if averageTrueRange > 0:
            positionSize = float(self.Portfolio.TotalPortfolioValue)  * self.riskFactor / float(averageTrueRange)
        return positionSize
        
    def SignificantChangeInPositionSize(context, newPositionSize, oldPositionSize):
        return np.abs((newPositionSize - oldPositionSize)  / oldPositionSize) > self.significantPositionDifference