from Alphas.RsiAlphaModel import RsiAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class MyAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.symbol = "SPY"
        self.res2use = Resolution.Daily
        
        self.SetStartDate(2018, 6, 17)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash

        # request the daily equity data
        self.AddEquity(self.symbol, self.res2use)
        self.AddPlots(self.symbol, self.res2use)
        
        # Six module plug and play algorithm development model
        self.AddAlpha(RsiAlphaModel(60, self.res2use))
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))
        # self.SetUniverseSelection(QC500UniverseSelectionModel())
        symbols = [ Symbol.Create(self.symbol, SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        if not self.Portfolio.Invested:
            self.SetHoldings(self.symbol, 1.0)
            
    def AddPlots(self, symbol, res2use):
        # Calcualte and plot various technical indicators
        self.sym_price = self.Identity(symbol)
        
        # Process: 1. Create Indictor 
        #          2. Register the daily data of "SPY" to automatically update the indicator 
        #          3. Plot indicator
        
        # SMA - Simple moving average
        self.sma50 = self.SMA(symbol, 50, res2use)
        self.sma200 = self.SMA(symbol, 200, res2use)
        self.RegisterIndicator(symbol, self.sma50)
        self.RegisterIndicator(symbol, self.sma200)
        self.PlotIndicator("SMA50-SMA200", self.sym_price, self.sma50, self.sma200)

        # BB - Bolling Bands 
        self.bb = self.BB(symbol, 200, res2use)
        self.RegisterIndicator(symbol, self.bb)
        self.PlotIndicator("BB", self.sym_price, self.bb.UpperBand, self.bb.LowerBand)

        # RSI - Relative Strength Index
        self.rsi = self.RSI(symbol, 10, MovingAverageType.Simple, res2use)
        self.RegisterIndicator(symbol, self.rsi)
        self.PlotIndicator("RSI", self.rsi)