#region imports
from AlgorithmImports import *
#endregion

class FundamentalFactorAlphaModel(AlphaModel):
   
    def __init__(self):
        self.rebalanceTime = datetime.min
        # Dictionary containing set of securities in each sector
        # e.g. {technology: set(AAPL, TSLA, ...), healthcare: set(XYZ, ABC, ...), ... }
        self.sectors = {}
    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'''
        if algorithm.Time <= self.rebalanceTime:
            return []
        
        # Set the rebalance time to match the insight expiry
        self.rebalanceTime = Expiry.EndOfQuarter(algorithm.Time)
        
        insights = []
        
        for sector in self.sectors:
            securities = self.sectors[sector]
            sortedByROE_pos = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=True)
            sortedByPM_pos = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=True)
            sortedByPE_pos = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=True)
            sortedByROE_neg = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=False)
            sortedByPM_neg = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=False)
            sortedByPE_neg = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=False)
            # Dictionary holding a dictionary of scores for each security in the sector
            scores_pos = {}
            for security in securities:
                score = sum([sortedByROE_pos.index(security), sortedByPM_pos.index(security), sortedByPE_pos.index(security)])
                scores_pos[security] = score

            scores_neg = {}
            for security in securities:
                score = sum([sortedByROE_neg.index(security), sortedByPM_neg.index(security), sortedByPE_neg.index(security)])
                scores_neg[security] = score
                
            # Add best 20% of each sector to longs set (minimum 1)
            length = max(int(len(scores_pos)/5), 1)
            for security in sorted(scores_pos.items(), key=lambda x: x[1], reverse=False)[:length]:
                symbol = security[0].Symbol
                # Use Expiry.EndOfQuarter in this case to match Universe, Alpha and PCM
                insights.append(Insight.Price(symbol, Expiry.EndOfQuarter, InsightDirection.Up))

            length = max(int(len(scores_neg)/5), 1)
            for security in sorted(scores_neg.items(), key=lambda x: x[1], reverse=False)[:length]:
                symbol = security[0].Symbol
                # Use Expiry.EndOfQuarter in this case to match Universe, Alpha and PCM
                insights.append(Insight.Price(symbol, Expiry.EndOfQuarter, InsightDirection.Down))
        
        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:
            sector = security.Fundamentals.AssetClassification.MorningstarSectorCode
            if sector not in self.sectors:
                self.sectors[sector] = set()
            self.sectors[sector].add(security)

# region imports
from AlgorithmImports import *
# endregion

from AlphaModel import *
class VerticalTachyonRegulators(QCAlgorithm):
   def Initialize(self):
       self.SetStartDate(2020, 1, 1)
       self.SetEndDate(2021, 1, 1)
       self.SetCash(100000)
       # Universe selection
       self.month = 0
       self.num_coarse = 500
       self.UniverseSettings.Resolution = Resolution.Daily
       self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
       
       # Alpha Model
       self.AddAlpha(FundamentalFactorAlphaModel())
       # Portfolio construction model
       self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(self.IsRebalanceDue))
       
       # Risk model
       self.SetRiskManagement(NullRiskManagementModel())
       # Execution model
       self.SetExecution(ImmediateExecutionModel())
   # Share the same rebalance function for Universe and PCM for clarity
   def IsRebalanceDue(self, time):
       # Rebalance on the first day of the Quarter
       if time.month == self.month or time.month not in [1, 4, 7, 10]:
           return None
           
       self.month = time.month
       return time
   def CoarseSelectionFunction(self, coarse):
       # If not time to rebalance, keep the same universe
       if not self.IsRebalanceDue(self.Time): 
           return Universe.Unchanged
       # 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 selected sectors
       
       sectors = [
           MorningstarSectorCode.FinancialServices,
           MorningstarSectorCode.RealEstate,
           MorningstarSectorCode.Healthcare,
           MorningstarSectorCode.Utilities,
           MorningstarSectorCode.Technology]
       
       filtered_fine = [x.Symbol for x in fine if x.SecurityReference.IPODate + timedelta(365*5) < self.Time
                                   and x.AssetClassification.MorningstarSectorCode in sectors]
               
       return filtered_fine