from HistoricalReturnsAlphaModel import HistoricalReturnsAlphaModel

from Risk.NullRiskManagementModel import NullRiskManagementModel


class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.coarse_count = 5
        self.averages = {}

        self.SetStartDate(2018, 6, 1)   #Set Start Date
        self.SetEndDate(2018, 9, 29)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        self.SetUniverseSelection(CoarseFundamentalUniverseSelectionModel(self.CoarseSelectionFunction))
        
        self.SetAlpha(HistoricalReturnsAlphaModel(14, Resolution.Daily))
        
        self.SetPortfolioConstruction(NullPortfolioConstructionModel()) 
        
        self.SetExecution(NullExecutionModel())
        
        self.SetRiskManagement(NullRiskManagementModel())
        
    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):

        # We are going to use a dictionary to refer the object that will keep the moving averages
        for cf in coarse:
            if cf.Symbol not in self.averages:
                self.averages[cf.Symbol] = SymbolData(cf.Symbol)

            # Updates the SymbolData object with current EOD price
            avg = self.averages[cf.Symbol]
            avg.update(cf)

        # Filter the values of the dict: wait for indicator to be ready
        filtered_values = filter(lambda x: (x.is_ready and x.price.Current.Value > 20), self.averages.values())

        # Sorts the values of the dict: we want those with greater DollarVolume
        sorted_values = sorted(filtered_values, key = lambda x: x.vol.Current.Value, reverse = True)

        for x in sorted_values[:self.coarse_count]:
            self.Log('symbol: ' + str(x.symbol.Value) + ' close price: ' + str(x.price.Current.Value) + '  mean vol: ' + str(x.vol.Current.Value) + '  mean price: ' + str(x.sma.Current.Value))
        
        # we need to return only the symbol objects
        return [ x.symbol for x in sorted_values[:self.coarse_count] ]

def OnOrderEvent(self, orderEvent):
    if orderEvent.Status == OrderStatus.Filled:
        # self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))
        pass

class SymbolData(object):
    
    def __init__(self, symbol):
        self.symbol = symbol
        self.price = SimpleMovingAverage(1)
        self.vol = SimpleMovingAverage(50)
        self.sma = SimpleMovingAverage(20)
        self.is_ready = False

    def update(self, value):
        self.is_ready = self.sma.Update(value.EndTime, value.Price) and self.vol.Update(value.EndTime, value.DollarVolume) and self.price.Update(value.EndTime, value.Price)

from clr import AddReference
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework.Alphas import *
from datetime import timedelta

class HistoricalReturnsAlphaModel(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 > 0: direction = InsightDirection.Up
                if magnitude < 0: direction = InsightDirection.Down

                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.ROC = RateOfChange('{}.ROC({})'.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.ROC, 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.ROC.Update(tuple.Index, tuple.close)

    @property
    def Return(self):
        return float(self.ROC.Current.Value)

    @property
    def CanEmit(self):
        if self.previous == self.ROC.Samples:
            return False

        self.previous = self.ROC.Samples
        return self.ROC.IsReady

    def __str__(self, **kwargs):
        return '{}: {:.2%}'.format(self.ROC.Name, (1 + self.Return)**252 - 1)