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

class QuantumTransdimensionalComputer(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2003, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 12, 31)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        # Momentum parameters
        self.mom = {}           # Dict of Momentum indicator keyed by Symbol
        self.lookback = 90     # Momentum indicator lookback period
        self.num_long = 30       # Number of symbols with open positions
        # Trade control Parameters
        self.month = -1     # month indicator
        self.rebalance = False # rebalance indicator


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        # Update the indicator
        for symbol, mom in self.mom.items():
            # Update also works for an IndicatorDataPoint object
            mom.Update(self.Time, self.Securities[symbol].Close)

        if not self.rebalance:
            return

        # Selects the securities with highest momentum
        # Note that sorted_mom and selected are lists of symbols
        sorted_mom = sorted([k for k,v in self.mom.items() if v.IsReady],
            key=lambda x: self.mom[x].Current.Value, reverse=True)
        selected = sorted_mom[:self.num_long]

        # Liquidate securities that are not in the list
        for symbol, mom in self.mom.items():
            if symbol not in selected:
                self.Liquidate(symbol, 'Not selected')

        # Buy selected securities
        for symbol in selected:
            self.SetHoldings(symbol, 1/self.num_long)

        self.rebalance = False


    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
        # Only one rebalance per mounth
        if self.month == self.Time.month:
            return Universe.Unchanged
        # Set rebalance indicator and reset month variable
        self.rebalance = True
        self.month = self.Time.month

        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5],
            key=lambda x: x.DollarVolume, reverse=True)

        return [x.Symbol for x in selected]
        
    
    def FineSelectionFunction(self, fine):
        '''Select security with highest market cap'''

        fine = [f for f in fine if f.ValuationRatios.PERatio > 0
                               and f.EarningReports.BasicEPS.TwelveMonths > 0
                               and f.EarningReports.BasicAverageShares.ThreeMonths > 0]
        # This is how Market Cap is calculated
        selected = [ f for f in fine if f.ValuationRatios.PERatio *
                                      f.EarningReports.BasicEPS.TwelveMonths *
                                      f.EarningReports.BasicAverageShares.ThreeMonths > 8200000000]
        return [x.Symbol for x in selected]
        
        
    def OnSecuritiesChanged(self, changes):
        # Clear removed securities
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if self.mom.pop(symbol, None) is not None:
                self.Liquidate(symbol, 'Removed from universe')
                
        # Create indicators for the new secuirities (NOT UPDATING YET)
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.mom:
                self.mom[symbol] = Momentum(self.lookback)
                
        # Warm up the indicator with history price if it is not ready
        addedSymbols = [k for k,v in self.mom.items() if not v.IsReady]
        # Get historical prices for the newly added securities
        # Use 1 + self.lookback as the window to be safer
        history = self.History(addedSymbols, 1 + self.lookback, Resolution.Daily)
        # The unstack changes a multiindex data frame to a single index data frame
        # Level = 0 keeps inner index and create columns based on the outer index
        history = history.close.unstack(level=0)
        # Manually update the indicator for the new securiites
        for symbol in addedSymbols:
            ticker = str(symbol)
            if ticker in history:
                # history[ticker] selects a *column*(result of unstack) based on the ticker
                for time, value in history[ticker].items():
                    # Manually create an indicator data point object and update with it
                    item = IndicatorDataPoint(symbol, time, value)
                    self.mom[symbol].Update(item)