from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class NadionUncoupledPrism(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(2010, 1, 1)
        self.SetEndDate(2011, 1, 1)
        self.SetCash(100000)
        
        #self.AddEquity("SPY", Resolution.Daily)
        #self.AddEquity("QQQ", Resolution.Daily)
        self.SetBenchmark("SPY")
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.symbols = [Symbol.Create("SPY", SecurityType.Equity, Market.USA) for ticker in ['SPY', 'TLT', 'QQQ']]
        self.averages = {}

        self.AddUniverseSelection(TechnologyUniverseModule())
        self.AddRiskManagement(NullRiskManagementModel())
        self.SetPortfolioConstruction(NullPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())

        self.SetWarmUp(200)

    def OnData(self, data):

        if self.IsWarmingUp:
            return

        #self.MarketOrder('SPY', 1)
        for security in self.changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)

        for security in self.changes.AddedSecurities:
            if not security.Invested and security.Symbol not in self.symbols:
                self.SetHoldings(security.Symbol, .05)

    def OnSecuritiesChanged(self, changes):
        self.changes = changes


class TechnologyUniverseModule(FundamentalUniverseSelectionModel):
    
    #This module selects the most liquid stocks listed on the Nasdaq Stock Exchange.
    
    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        #Initializes a new default instance of the TechnologyUniverseModule
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 100
        self.numberOfSymbolsFine = 100
        self.lastMonth = -1
        self.averages = {}


    def SelectCoarse(self, algorithm, coarse):
        if algorithm.Time.month == self.lastMonth: 
            return Universe.Unchanged
            
        filtered = [x for x in coarse if x.HasFundamentalData and x.Price > 10]
        if len(filtered) == 0:
            return Universe.Unchanged
        sorted_by_dollar_volume = sorted(filtered, key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse] 
        
        self.last_coarse = {cf.Symbol: cf for cf in sorted_by_dollar_volume}
        return list(self.last_coarse.keys())


    def SelectFine(self, algorithm, fine):
        
        filtered =  [x for x in fine if x.CompanyReference.CountryId == "USA" \
                                        and x.CompanyReference.PrimaryExchangeID == "NAS" \
                                        and x.CompanyReference.IndustryTemplateCode == "N" \
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180]

        sorted_by_dollar_volume = sorted(filtered, key = lambda x: self.last_coarse[x.Symbol].DollarVolume, reverse=True)
        
        # Get warmup history
        new_symbols = []
        for security in sorted_by_dollar_volume:
            symbol = security.Symbol
            if symbol not in self.averages:
                new_symbols.append(symbol)
        history = algorithm.History(new_symbols, 200, Resolution.Daily)
        
        selected = []
        for security in sorted_by_dollar_volume:
            symbol = security.Symbol
            if symbol in new_symbols:
                self.averages[symbol] = SelectionData(history.loc[symbol]) 
            self.averages[symbol].update(algorithm.Time, self.last_coarse[symbol].AdjustedPrice) #security.AdjustedPrice)
            if self.averages[symbol].is_ready() and self.averages[symbol].fast > self.averages[symbol].slow:
                selected.append(symbol)

        if len(selected) == 0:
            return Universe.Unchanged
        self.lastMonth = algorithm.Time.month
        return selected[:20]
        
        
class SelectionData():
    
    # Update the constructor to accept a history array
    def __init__(self, history):
        
        self.slow = ExponentialMovingAverage(200)
        self.fast = ExponentialMovingAverage(5)
        
        # Loop over the history data and update the indicators
        for time, row in history.iterrows():
            
            self.fast.Update(time, row.close)
            self.slow.Update(time, row.close)
    
    
    def is_ready(self):
        
        return self.slow.IsReady and self.fast.IsReady
    
    
    def update(self, time, price):
        
        self.fast.Update(time, price)
        self.slow.Update(time, price)
        
        
        
class NullPortfolioConstructionModel(PortfolioConstructionModel):

    def CreateTargets(self, algorithm, insights):
        
        return []
        
        
class ImmediateExecutionModel(ExecutionModel):

    def __init__(self):
        
        self.targetsCollection = PortfolioTargetCollection()
        

    def Execute(self, algorithm, targets):

        # for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
        self.targetsCollection.AddRange(targets)
        
        if self.targetsCollection.Count > 0:
            
            for target in self.targetsCollection.OrderByMarginImpact(algorithm):
                # calculate remaining quantity to be ordered
                
                quantity = OrderSizing.GetUnorderedQuantity(algorithm, target)
                
                if quantity != 0:
                    
                    algorithm.MarketOrder(target.Symbol, quantity)


            self.targetsCollection.ClearFulfilled(algorithm)

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class TechnologyUniverseModule(FundamentalUniverseSelectionModel):
    '''
    This module selects the most liquid stocks listed on the Nasdaq Stock Exchange.
    '''
    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the TechnologyUniverseModule'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = 100
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1

    def SelectCoarse(self, algorithm, coarse):
        '''
        Performs a coarse selection:
        
        -The stock must have fundamental data
        -The stock must have positive previous-day close price
        -The stock must have positive volume on the previous trading day
        '''
        if algorithm.Time.month == self.lastMonth: 
            return Universe.Unchanged

        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
            key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}
        
        # If no security has met the QC500 criteria, the universe is unchanged.
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''
        Performs a fine selection:
        
        -The company's headquarter must in the U.S.
        -The stock must be traded on the NASDAQ stock exchange
        -The stock must be in the Industry Template Code catagory N
        -At least half a year since its initial public offering
        '''
        # Filter stocks and sort on dollar volume
        sortedByDollarVolume = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID == "NAS"
                                        and x.CompanyReference.IndustryTemplateCode == "N"
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180],
            key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)

        if len(sortedByDollarVolume) == 0:
            return Universe.Unchanged
            
        self.lastMonth = algorithm.Time.month

        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]