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

from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Indicators import ExponentialMovingAverage
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

import decimal as d

class UniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self,
                 universeCount = 20,
                 universeSettings = None,
                 securityInitializer = None):

        super().__init__(False, universeSettings, securityInitializer)
        self.universeCount = universeCount
        self.averages = {}

    def SelectCoarse(self, algorithm, 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 price and volume
            avg = self.averages[cf.Symbol]
            avg.update(cf.EndTime, cf.Price, cf.Volume)
        
        # Filter the values of the dict based on criteria
        # In this case, we only want up-trending securities
        filtered_values = list(filter(lambda x: x.is_uptrend, self.averages.values()))
        
        # Sort the values of the dict: we want those with greater DollarVolume
        sorted_values = sorted(filtered_values, key = lambda x: x.volume, reverse = True)
        
        return [ x.symbol for x in sorted_values[:self.universeCount] ]

# class used to improve readability of the coarse selection function
class SymbolData:
    def __init__(self, symbol):
        self.symbol = symbol
        self.priceSMA = SimpleMovingAverage(10)
        self.volSMA = SimpleMovingAverage(50)
        self.priceWin = RollingWindow[float](5)
        self.is_uptrend = False

    def update(self, time, price, volume):
        self.price = price
        self.volume = volume
        self.priceWin.Add(price)
        self.priceSMA.Update(time, price)
        self.volSMA.Update(time, volume)
        if self.priceSMA.IsReady and self.volSMA.IsReady and self.priceWin.IsReady:
            MAprice = self.priceSMA.Current.Value
            MAvol = self.volSMA.Current.Value
            
            # Here we can add the criteria for our universe
            # current price > 10-days moving average price
            # current volume > 10-days moving average volume
            # current price > yesterday's close
            # current price > 10-days maximum price
            
            self.is_uptrend =  price > 20 and price > MAprice and volume > MAvol and price > self.priceWin[1] and price > max(self.priceWin)

from UniverseSelection import UniverseSelectionModel
from Alphas.ConstantAlphaModel import ConstantAlphaModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Risk.NullRiskManagementModel import NullRiskManagementModel

class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Daily

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

        count = 5
        self.SetUniverseSelection(UniverseSelectionModel(universeCount = count))
        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, TimeSpan.FromMinutes(20), 0.025, None))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        
    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))
            
        #for x in sorted_values[:self.universeCount]:
            #self.Log('symbol: ' + str(x.symbol.Value)
            #+ ' close price: ' + str(x.price)
            #+ ' volume: ' + str(x.volume)
            #+ '  mean price: ' + str(x.priceSMA.Current.Value)
            #+ '  mean vol: ' + str(x.volSMA.Current.Value))
        
    def OnEndOfDay(self):
        
        self.Liquidate()

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)