from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class CalibratedOptimizedCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
    
        # Set QC500 Universe Selection Model
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Quality Companies Alpha 
        self.SetAlpha(QualityCompaniesAlpha())
            
        # Equally weigh securities in portfolio, based on insights
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        # Set Immediate Execution Model
        self.SetExecution(ImmediateExecutionModel())

        # Set Null Risk Management Model
        self.SetRiskManagement(NullRiskManagementModel())
        
class QualityCompaniesAlpha(AlphaModel): 
    
    def __init__(self):
        self.SecData = {}
        self.selected = []
        self.month = -1
        
    def Update(self, algorithm, data):
        # initialize insights 
        insights = []
        
        # only trade once a month a few days before months end
        d1 = (algorithm.Time.day != 24)
        d2 = (algorithm.Time.day != 25) 
        d3 = (algorithm.Time.day != 26)
        if d1 or d2 or d3: 
            return insights
        
        if self.month == algorithm.Time.month: 
            return insights
        
        self.month = algorithm.Time.month
        
        # SPY change should be positive 
        spyhist = algorithm.History(spy, 140, Resolution.Daily)
        spychange = spyhist.close.pct_change(periods=139)[-1]
        
        if spychange > 0: 
            self.selected = self.SelectStocks(data)
        else: 
            return insights 
        
        for kvp in algorithm.Portfolio:
            symbol = kvp.Key
            holding = kvp.Value
            if symbol not in self.selected and holding.Invested:
                insights.append(Insight.Price(symbol, timedelta(30), InsightType.Price, InsightDirection.Flat, None, None))
        
        for symbol in self.selected:
            if not algorithm.Portfolio[symbol].Invested:
                insights.append(Insight.Price(symbol, timedelta(30), InsightType.Price, InsightDirection.Up, None, None))
        
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.SecData:
                # adds filter criteria
                self.SecData[symbol] = SecurityData(algorithm, symbol, security)
                
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.SecData:
                # Remove consolidator for removed securities
                algorithm.SubscriptionManager.RemoveConsolidator(symbol, self.SecData[symbol].consolidator)
                self.SecData.pop(symbol, None)
                
    def SelectStocks(self, data):
        # Sort by ROIC
        sortedByROIC = sorted(list(self.SecData.keys()), key=lambda k : self.SecData[k].roic, reverse=True)
        
        # Sorted by MOM
        sortedByMOM = sorted(sortedByROIC[:50], key=lambda k : self.SecData[k].mom.Current.Value, reverse=True)
        
        # Select top 20 symbols
        return sortedByMOM[:20]
        
class SecurityData: 
    def __init__(self, algorithm, symbol, security): 
        # general 
        self.algorithm = algorithm
        self.symbol = symbol
        self.security = security
        # adds filter criteria
        self.roic = security.Fundamentals.OperationRatios.ROIC.OneYear
        
        # Define consolidator for updating mom
        self.consolidator = TradeBarConsolidator(timedelta(days = 1))
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)
        
        # Define and register mom indicator
        self.mom = Momentum(126)
        self.algorithm.RegisterIndicator(symbol, self.mom, self.consolidator)
        
        # Initialize MOM indicator with historical data
        history = algorithm.History(symbol, 126, Resolution.Daily)
        if not history.empty:
            history = history.close.unstack(0)
            if not history.empty:
                df = history[symbol]
                for time, close in df.iteritems():
                    self.mom.Update(time, close)