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


class CapitalizationUniverseModel(FundamentalUniverseSelectionModel):
    """
    Universe model that selects stocks by capitalization. It performs an initial
    selection by traded volume to improve algorithm speed.
    """

    def __init__(self, filterFineData=True, universeSettings=None, securityInitializer=None):

        '''
        Initializes a new default instance of the Capitalization Universe Model.
        '''

        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 2000
        self.num_symbols = 50
        self.price_limit = 1
        self.last_month = -1 # Start first period
        self.is_take_top = False

    def SelectCoarse(self, algorithm, coarse):
        
        '''
        Drop securities which have no fundamental data or have too low prices.
        Select those by dollar volume ordered accorsing to self.is_take_top.
        '''
    
        if algorithm.Time.month == self.last_month:
            return Universe.Unchanged
        self.last_month = algorithm.Time.month
        
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > self.price_limit],
                            key=lambda x: x.DollarVolume, reverse=self.is_take_top)
    
        return [x.Symbol for x in selected[:self.numberOfSymbolsCoarse]]

    def SelectFine(self, algorithm, fine):
        ''' Selects the stocks by market cap '''
        sorted_market_cap = sorted([x for x in fine if x.MarketCap > 0],
        key=lambda x: x.MarketCap, reverse = self.is_take_top)

        return [x.Symbol for x in sorted_market_cap[:self.num_symbols]]

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from datetime import datetime, timedelta
from CapitalizationUniverseModel import CapitalizationUniverseModel
from QuantConnect.Indicators import Momentum


class MomentumLaboratory(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(datetime.today())
        self.SetCash(1000000)
        self.AddEquity("SPY", Resolution.Daily)
        self.SetBenchmark("SPY")
        self.SetUniverseSelection(CapitalizationUniverseModel())
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.TotalReturn))
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.AddAlpha(MomentumReversalAlphaModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())


class MomentumReversalAlphaModel(AlphaModel):
    """
    Demostration Momentum Indicator Alpha Model.
    """

    def __init__(self):

        self.Name = 'MomentumReversalAlphaModel'
        self.is_reversal = True
        # Re-balance every week.
        self.rebalancing_period = 30
        # Initiate first day of operation.
        self.next_rebalance = False
        # EMA period, in days.
        self.indicator_period = 30
        # Gain Target:
        self.gain_target = 0.05
        self.max_positions = 15

        # Dictionary to hold symbol data:
        self.symbol_data = {}

        # Flag to show first day of work completed.
        self.start_flag = False
        self.directions = [InsightDirection.Down, InsightDirection.Up]

    def Update(self, algorithm, data):

        # Clear counters and insights, return if data slice is empty.
        insights = []
        up_count = False
        if not data.HasData:
            return []

        # If this is the first day of operation, lower the flag and set
        # next re-balance time for today.
        if self.start_flag is False:
            self.next_rebalance = algorithm.Time
            self.start_flag = True

        if algorithm.Time < self.next_rebalance:
            return[]

        # Obtain a list of securities both in the indicator dictionary and in
        # the data slice.
        common_list = list(set(data.Keys).intersection(self.symbol_data.keys()))
        ordered_list = sorted(common_list, key=lambda x: self.symbol_data[x].indicator.Current.Value, reverse=True)
        up_momentum = [symbol for symbol in ordered_list][:self.max_positions]
        down_momentum = [symbol for symbol in ordered_list][-self.max_positions:]

        # Emit the insight for rebalance period.
        for symbol in up_momentum + down_momentum:
            # Avoid SPY
            if symbol == algorithm.Symbol("SPY"):
                continue

            if data[symbol]:
                if symbol in up_momentum:
                    direction = True
                if symbol in down_momentum:
                    direction = False
                if self.is_reversal:
                    direction = not direction
                direction = self.directions[direction]

                insight = Insight(symbol,
                                  timedelta(days=self.rebalancing_period),
                                  InsightType.Price,
                                  direction, self.gain_target, 1, self.Name, 1)
                insights.append(insight)

        # Set next re-balancing date.
        self.next_rebalance = algorithm.Time + timedelta(days=self.rebalancing_period)

        return Insight.Group(insights)

    def OnSecuritiesChanged(self, algorithm, changes):

        # Generate and warm-up indicator for added securities:
        for security in changes.AddedSecurities:
            if security.Symbol not in self.symbol_data:
                self.symbol_data[security.Symbol] = SymbolData(security.Symbol,
                                                               algorithm,
                                                               self.indicator_period)

        for symbol in self.symbol_data.keys():
            history = algorithm.History(security.Symbol,
                                        self.indicator_period,
                                        Resolution.Daily)
            if not history.empty:
                self.symbol_data[security.Symbol].WarmUpIndicators(history)

        for security in changes.RemovedSecurities:
            if security.Symbol in self.symbol_data:
                algorithm.SubscriptionManager.RemoveConsolidator(security.Symbol, self.symbol_data[security.Symbol].consolidator)
                self.symbol_data.pop(security.Symbol)
        return


# Class to hold required symbol data
# Adapted from EMACrossOverUniverseSelectionModel
class SymbolData(object):

    def __init__(self, symbol, algorithm, period):
        self.symbol = symbol
        self.indicator = Momentum(period)
        self.consolidator = algorithm.RegisterIndicator(self.symbol,
                                                        self.indicator,
                                                        Resolution.Daily)

    def WarmUpIndicators(self, history):
        for time, row in history.loc[self.symbol].iterrows():
                self.indicator.Update(time, row["close"])