from System import *
from clr import AddReference
AddReference("QuantConnect.Algorithm")
from QuantConnect import *
from QuantConnect.Orders import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Execution import ExecutionModel
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioConstructionModel
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class Algorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 4, 1)
        self.SetCash(100000)
        
        self.AddUniverseSelection(TechnologyUniverseModule())
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddEquity("SPY", Resolution.Daily)
        self.AddEquity("QQQ", Resolution.Daily)
        
        self.sma = self.SMA("SPY", 200, Resolution.Daily)

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.changes = None

        self.SetWarmUp(100)
        self.symbols = [Symbol.Create("QQQ", SecurityType.Equity, Market.USA), \
        Symbol.Create("SPY", SecurityType.Equity, Market.USA)]

    def OnData(self, data):
        if self.IsWarmingUp:
            return
        if not self.changes == None:
            for security in self.changes.AddedSecurities:
                self.SetHoldings(security.Symbol, 0.2)
        # Plot number of active securities
        numActiveSecurities = 0
        for security in self.ActiveSecurities.Values:
            numActiveSecurities += 1
        self.Plot("Active", "Securities", int(numActiveSecurities))
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol, 'Removed from Universe')
            # self.Log(f"OnSecuritiesChanged({self.UtcTime}):: {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 = 1000
        self.numberOfSymbolsFine = 100
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1
        

    def SelectCoarse(self, algorithm, coarse):
        '''
        Coarse Filters:
        -The stock must have fundamental data
        -The stock must have positive previous-month close price
        -The stock must have positive volume on the previous trading month
        '''
        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 > 1], \
                               key = lambda x: x.DollarVolume, reverse=True)


        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):

        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[:5]]