from clr import AddReference
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Indicators import *

from datetime import timedelta, datetime

class rollingWindowExampleFW(QCAlgorithmFramework):

    def Initialize(self):
        self.SetStartDate(2019, 2, 1)
        self.SetCash(100000)
    
        # Add SPY to Universe
        self.AddSecurity(SecurityType.Equity,"SPY", Resolution.Daily)

        # Use Manual Universe Selection Model
        self.SetUniverseSelection(ManualUniverseSelectionModel())
        
        # Select the Alpha Model
        self.SetAlpha(rollingWindowExampleAlphaModel())
        
        # Equally weigh securities
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

class rollingWindowExampleAlphaModel(AlphaModel):

    def __init__(self):
        # Set length of the insight period
        self.period = timedelta(days=1)
        
        # Set lookback period for EMA and Rolling Window
        self.lookback = 5
        
        # Set data resolution
        self.resolution = Resolution.Daily

    def Update(self, algorithm, data):
        insights = []
        if not data.ContainsKey(self.symbol): return insights
        
        # Retrive current price
        price = algorithm.Securities[self.symbol].Price
        
        # Retrieve current EMA value from rolling window
        currentEMA = self.emaWin[0].Value
        
        # Set magnitude prediction as the average percentage change of price
        magnitude = self.meanOfPriceChange.Current.Value
        
        # Emit an "Up" insight if current price exceeds EMA
        if price > currentEMA:
            insights.append(Insight.Price(self.symbol, self.period, InsightDirection.Up,magnitude))

        return insights
    
    # Ignites when securities are added/removed from universe
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbol = added.Symbol
            
            # Mean value of returns for magnitude prediction
            self.meanOfPriceChange = IndicatorExtensions.SMA(RateOfChangePercent(1),self.lookback)
            
            # Construct an EMA indicator and a rolling window for added securities
            algorithm.EMA(added.Symbol, self.lookback,self.resolution).Updated += self.emaUpdated
            self.emaWin = RollingWindow[IndicatorDataPoint](self.lookback)
            
    # Adds values to rolling window
    def emaUpdated(self, sender, updated):
        self.emaWin.Add(updated)