from AlphaModel import *

class VerticalTachyonRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 1)
        self.SetCash(100000)
        
        # Execution model
        self.SetExecution(ImmediateExecutionModel())
        
        # Portfolio construction model
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(rebalance=timedelta(weeks=13)))
        
        # Risk model
        self.SetRiskManagement(NullRiskManagementModel())
        
        # Universe selection
        self.num_coarse = 500
        self.rebalanceTime = datetime.min
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.sectors = set([MorningstarSectorCode.FinancialServices, MorningstarSectorCode.RealEstate, MorningstarSectorCode.Healthcare, MorningstarSectorCode.Utilities, MorningstarSectorCode.Technology])
        self.period = timedelta(weeks=13)
        
        # Alpha Model
        self.AddAlpha(FundamentalFactorAlphaModel(self.period, self.sectors))
    

    def CoarseSelectionFunction(self, coarse):
        # If not time to rebalance, keep the same universe
        if self.Time <= self.rebalanceTime: 
            return Universe.Unchanged
        
        self.rebalanceTime = self.Time + self.period
        
        # Select only those with fundamental data and a sufficiently large price
        # Sort by top dollar volume: most liquid to least liquid
        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[:self.num_coarse]]


    def FineSelectionFunction(self, fine):
        # Filter the fine data for equities that IPO'd more than 5 years ago in given sector
        filtered_fine = [x.Symbol for x in fine if x.SecurityReference.IPODate + timedelta(365*5) < self.Time
                                    and x.AssetClassification.MorningstarSectorCode  in self.sectors
                                    and x.OperationRatios.ROE.Value > 0
                                    and x.OperationRatios.NetMargin.Value > 0
                                    and x.ValuationRatios.PERatio > 0]
                
        return filtered_fine

from datetime import timedelta

class FundamentalFactorAlphaModel(AlphaModel):
    
    def __init__(self, period, sectors):
        self.rebalanceTime = datetime.min
        self.sectors = {}
        for sector in sectors:
            self.sectors[sector] = set()
        self.period = period


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            New insights'''
        
        # Return no insights if it's not time to rebalance
        if algorithm.Time <= self.rebalanceTime: 
            return []
        self.rebalanceTime = algorithm.Time + self.period
        
        # Dictionary holding a dictionary of scores for each sector
        sector_scores = {}
        
        for sector in self.sectors:
            securities = self.sectors[sector]
            sortedByROE = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=True)
            sortedByPM = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=True)
            sortedByPE = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=False)
            scores = {}
            for security in securities:
                score = sum([sortedByROE.index(security), sortedByPM.index(security), sortedByPE.index(security)])
                scores[security] = score
            sector_scores[sector] = scores
        
        # Set of symbols that we want to buy
        longs = set()
        
        for sector in self.sectors:
            # add best 20% of each sector to longs set (minimum 1)
            length = max(int(len(sector_scores[sector])/5), 1)
            for security in sorted(sector_scores[sector].items(), key=lambda x: x[1], reverse=False)[:length]:
                longs.add(security[0].Symbol)
        
        # Insights of the form: Insight(symbol, timedelta, type, direction, magnitude, confidence, sourceModel, weight)
        insights = []
        
        # Close old positions if they aren't in longs set
        for security in algorithm.Portfolio.Values:
            if security.Invested and security.Symbol not in longs:
                insights.append(Insight(security.Symbol, Expiry.EndOfDay, InsightType.Price, InsightDirection.Flat))
        
        # Emit buy insight for all symbols in longs set
        for symbol in longs:
            insights.append(Insight(symbol, self.period, InsightType.Price, InsightDirection.Up))
        
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        
        # Remove security from sector set
        for security in changes.RemovedSecurities:
            for sector in self.sectors:
                if security in self.sectors[sector]:
                    self.sectors[sector].remove(security)
        
        # Add security to corresponding sector set
        for security in changes.AddedSecurities:
            self.sectors[security.Fundamentals.AssetClassification.MorningstarSectorCode].add(security)