from Execution.ImmediateExecutionModel import ImmediateExecutionModel

class InterdayMeanReversion(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 12, 1)  #. Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddUniverse(self.Universe.Index.QC500)
        
        self.AddAlpha(ImrAlphaModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetBenchmark("SPY")
        

class ImrAlphaModel(AlphaModel):
    
    def __init__(self):
        self.momp = {}
        self.momp_period = 5
        self.effect_period = 5
        self.num_of_selected_long = 10
        self.num_of_selected_short = 10
    
    def OnSecuritiesChanged(self, algorithm, changes):

        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.momp[symbol] = algorithm.MOMP(symbol, self.momp_period, Resolution.Daily)

        for security in changes.RemovedSecurities:
            mompInd = self.momp.pop(security.Symbol, None)
            if mompInd is not None:
                algorithm.RemoveSecurity(security.Symbol)
    
    def Update(self, algorithm, data):
        
        if algorithm.Time.weekday() != 0:
            return []
        
        insights = []
        insight_magnitude = {symbol: -1 * fiveDayPctMomentum.Current.Value for (symbol, fiveDayPctMomentum) in self.momp.items()}
        sorted_insight = sorted(insight_magnitude.items(), key = lambda item: item[1], reverse=True)
        
        top_pos_insight = sorted_insight[:self.num_of_selected_long]
        top_neg_insight = sorted_insight[-self.num_of_selected_short:]
        
        for symbol, signal in top_pos_insight:
            insights.append(Insight.Price(symbol, timedelta(self.effect_period), InsightDirection.Up))
            
        for symbol, signal in top_neg_insight:
            insights.append(Insight.Price(symbol, timedelta(self.effect_period), InsightDirection.Down))
        
        return Insight.Group(insights)