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.removedSecurities = 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
        numActiveSecurities = len(self.ActiveSecurities)
        if self.changes and numActiveSecurities > 0:
            for security in self.ActiveSecurities.Values:
                if not security in self.removedSecurities:
                    self.Log("Setting Holdings " + str(security.Symbol))
                    self.SetHoldings(security.Symbol, 1/numActiveSecurities, True)
            self.changes = None
        # Plot the number of active securities
        self.Plot("Active", "Securities", numActiveSecurities)
        
    def OnSecuritiesChanged(self, changes):
        self.changes = changes
        self.removedSecurities = changes.RemovedSecurities
        self.Log(f"OnSecuritiesChanged({self.UtcTime}) :: {changes}")
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol, 'Removed from Universe')

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]]