class MyHistoricalReturnsAlphaModel(AlphaModel):
    '''Uses Historical returns to create insights.'''
        
    def __init__(self, *args, **kwargs):
        '''Initializes a new default instance of the HistoricalReturnsAlphaModel class.
        Args:
            lookback(int): Historical return lookback period
            resolution: The resolution of historical data'''
        self.lookback = kwargs['lookback'] if 'lookback' in kwargs else 1
        self.resolution = kwargs['resolution'] if 'resolution' in kwargs else Resolution.Daily
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), self.lookback)
        self.symbolDataBySymbol = {}

    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.CanEmit:
                
                direction = InsightDirection.Flat
                magnitude = symbolData.Return
                if magnitude > .10: direction = InsightDirection.Down
                if magnitude < -.10: direction = InsightDirection.Up
                
                insights.append(Insight.Price(symbol, self.predictionInterval, direction, magnitude, None))
                
        return insights
            
    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
            
        # clean up data for removed securities
        for removed in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if symbolData is not None:
                symbolData.RemoveConsolidators(algorithm)
                
        # initialize data for added securities
        symbols = [ x.Symbol for x in changes.AddedSecurities ]
        history = algorithm.History(symbols, self.lookback, self.resolution)
        if history.empty: return
            
        tickers = history.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.GetSymbol(ticker)
                
            if symbol not in self.symbolDataBySymbol:
                symbolData = SymbolData(symbol, self.lookback)
                self.symbolDataBySymbol[symbol] = symbolData
                symbolData.RegisterIndicators(algorithm, self.resolution)
                symbolData.WarmUpIndicators(history.loc[ticker])
                
                
class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, symbol, lookback):
        self.Symbol = symbol
        self.ROCP = RateOfChange('{}.ROCP({})'.format(symbol, lookback), lookback)
        
        self.Consolidator = None
        self.previous = 0
            
    def RegisterIndicators(self, algorithm, resolution):
        self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        algorithm.RegisterIndicator(self.Symbol, self.ROCP, self.Consolidator)
            
    def RemoveConsolidators(self, algorithm):
        if self.Consolidator is not None:
            algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.Consolidator)
            
    def WarmUpIndicators(self, history):
        for tuple in history.itertuples():
            self.ROCP.Update(tuple.Index, tuple.close)
            
    @property
    def Return(self):
        return float(self.ROCP.Current.Value)
            
    @property
    def CanEmit(self):
        if self.previous == self.ROCP.Samples:
            return False
            
        self.previous = self.ROCP.Samples
        return self.ROCP.IsReady
        
    def __str__(self, **kwargs):
        return '{}: {:.2%}'.format(self.ROC.Name, (1 + self.Return)**252 - 1)

class BracketRiskModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the maximum possible loss
    measured from the highest unrealized profit'''
    def __init__(self, maximumDrawdownPercent = 0.05, maximumUnrealizedProfitPercent = 0.05):
        '''Initializes a new instance of the TrailingStopRiskManagementModel class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio compared with the highest unrealized profit, defaults to 5% drawdown'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.trailingHighs = dict()
        self.maximumUnrealizedProfitPercent = abs(maximumUnrealizedProfitPercent)
        
    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        riskAdjustedTargets = list()
        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            # Remove if not invested
            if not security.Invested:
                self.trailingHighs.pop(symbol, None)
                continue
            pnl = security.Holdings.UnrealizedProfitPercent
            
            if pnl > self.maximumUnrealizedProfitPercent:
                # liquidate
                algorithm.Debug("Profit Taken")
                algorithm.Log("Profit Taken")
                riskAdjustedTargets.append(PortfolioTarget(security.Symbol, 0))
                return riskAdjustedTargets
            # Add newly invested securities
            if symbol not in self.trailingHighs:
                self.trailingHighs[symbol] = security.Holdings.AveragePrice   # Set to average holding cost
                continue

            # Check for new highs and update - set to tradebar high
            if self.trailingHighs[symbol] < security.High:
                self.trailingHighs[symbol] = security.High
                continue

            # Check for securities past the drawdown limit
            securityHigh = self.trailingHighs[symbol]
            drawdown = (security.Low / securityHigh) - 1

                
            if drawdown < self.maximumDrawdownPercent:
                # liquidate
                algorithm.Debug("Losses Taken")
                algorithm.Log("Losses Taken")
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                
        return riskAdjustedTargets

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

from AlgorithmImports import *

from risk import BracketRiskModel
from alpha import MyHistoricalReturnsAlphaModel

class TwoWeekMagnitudeEffectOnCurrentReturns(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 5, 8)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.AddAlpha(MyHistoricalReturnsAlphaModel(14, Resolution.Daily))
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        # uses universe selection to build a universe with 150 high volume stocks that have earnings within 3 months and low ATR(Average True Range) which signals low volatility
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction)) 
        self.UniverseSettings.Resolution = Resolution.Daily # makes sure the universe is on a daily rebalance period
        self.period = 14
        
        self.up_value = 0.10
        self.down_value = 0.05
        self.SetRiskManagement(BracketRiskModel(self.down_value, self.up_value))
        
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        
    def CoarseSelectionFunction(self, coarse):
        # sort the universe by high dollar volume
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        # filter any assets that are above $5.00 and have a trading volume of over 10,000,000. then return the top 275
        filtered = [ x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData and x.Price >= 5 and x.DollarVolume > 10000000 ][:275]
        return [ x for x in filtered ] # passes to FineSelectionFunction()
        
    def FineSelectionFunction(self, fine):
        # filters assets to return the top 150 that have earnings dates within 3 months
        filteredByDates = [ x for x in fine if x.EarningReports.FileDate.year == self.Time.year and x.EarningReports.FileDate.month - self.Time.month <= 3 ][:150]
        
        # gives us our atr indicator values
        for symbol in filteredByDates:
            self.AddSecurity(symbol.Symbol, Resolution.Daily)
            history = self.History([symbol.Symbol], self.period, Resolution.Daily)
            atr = self.ATR(symbol.Symbol, self.period, Resolution.Daily)
            for bar in history.itertuples():
                tradebar = TradeBar(bar.Index[1], symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
                atr.Update(tradebar)
            symbol.atr = atr.Current.Value
        
        # use this to sort by low to high ATR values 
        sortedByVolatility = sorted(filteredByDates, key=lambda x: abs(x.atr), reverse=False)
        
        return [ x.Symbol for x in sortedByVolatility ] # gives us our universe