# Import packages
import numpy as np
import pandas as pd
import scipy as sc


#class InOut(QCAlgorithm):
#
#    def Initialize(self):
#
#        self.SetStartDate(2021, 1, 1)  #Set Start Date
#        self.SetEndDate(2021, 4, 26)  #Set End Date
#
#        
#        
#        
#    def CoarseSelectionFunction(self, coarse):
#        if not self.selection_flag:
#            return Universe.Unchanged
#        
#        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Market == 'usa' and x.Price > 5],
#        key=lambda x: x.DollarVolume, reverse=True)
#        
#        return [x.Symbol for x in selected[:self.course_count]]
        
        
        
class NadionResistanceShield(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 1, 1)
        self.SetCash(25000)  # Set Strategy Cash
        self.tickers =  [ "TSLA"] 
        self.symbolDataBySymbol = {}
        self.trade = True
        self.atr=[]
        
        self.MarketCaps = ["IWM", "MDY", "SPY", "QQQ"]  
        self.marketDataBySymbol = {}
        
        # Before the open
            
        
        
        for ticker_mark in self.MarketCaps:
            symbol = self.AddEquity(ticker_mark, Resolution.Daily).Symbol
            
            sma50 = self.SMA(symbol, 50, Resolution.Daily, Field.Close)
            sma200 = self.SMA(symbol, 200, Resolution.Daily, Field.Close)
        
            self.marketDataBySymbol[symbol] = symbolMarkData(symbol, sma50, sma200)
        
        
        
        
        for ticker in self.tickers:
            symbol = self.AddEquity(ticker, Resolution.Hour).Symbol
            
            '''For the below 3 EMA's, you can convert them to 4H bars using the colidator method'''
            
            sma10 = self.SMA(symbol, 10, Resolution.Daily, Field.Close)
            sma200 = self.SMA(symbol, 200, Resolution.Daily, Field.Close)
            sma7 = self.SMA(symbol, 7, Resolution.Hour, Field.Close)
            sma20 = self.SMA(symbol, 20, Resolution.Daily, Field.Close)
            sma50 = self.SMA(symbol, 50, Resolution.Daily, Field.Close)
            sma250 = self.SMA(symbol, 250, Resolution.Daily, Field.Close)
            ema20 = self.EMA(symbol, 20, Resolution.Hour, Field.Close)
            ema50 = self.EMA(symbol, 50, Resolution.Hour, Field.Close)
            rsi = self.RSI(symbol, 14, Resolution.Daily)
            wilr = self.WILR(symbol, 14, Resolution.Daily)
            wilr_fast = self.WILR(symbol, 10, Resolution.Daily)
            atr = self.ATR(symbol, 14, Resolution.Daily)
            
             
            self.atr.append(self.ATR(symbol, 7, Resolution.Daily))
        
            
            '''Consolidator method'''
            sma20_4hr = ExponentialMovingAverage(20, MovingAverageType.Simple)#, Resolution.Hour, Field.Close)
            sma250_4hr = ExponentialMovingAverage(250, MovingAverageType.Simple)
            # create the 4 hour data consolidator
            fourHourConsolidator = TradeBarConsolidator(timedelta(hours=4))
            self.SubscriptionManager.AddConsolidator(symbol, fourHourConsolidator)
            # register the 4 hour consolidated bar data to automatically update the indicator
            self.RegisterIndicator(symbol, sma20_4hr, fourHourConsolidator)
            self.RegisterIndicator(symbol, sma250_4hr, fourHourConsolidator)
            
            #self.Schedule.On(self.DateRules.EveryDay(self.tickers), 
            #self.TimeRules.AfterMarketOpen(self.tickers, -5), 
              #  Action(self.beforeTheOpen))
            
            symbolData = SymbolData(symbol, sma10, sma20, sma200, sma7, sma50, sma250, ema20, ema50, rsi, wilr, wilr_fast, atr, sma20_4hr, sma250_4hr)
            self.symbolDataBySymbol[symbol] = symbolData
        
        
        
        self.spy = self.AddEquity("SPY", Resolution.Daily)
        
        # Before the open
        self.Schedule.On(self.DateRules.EveryDay("SPY"), 
                self.TimeRules.AfterMarketOpen("SPY", -5), 
                Action(self.beforeTheOpen))
        
        #set the following between 1 - 4 hours depending on buy frequency    
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=1)),
                 self.buySignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=.5)),
                 self.sellSignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY"),
                 self.tradeStart)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY"),
                 self.tradeEnd)

        #self.AddRiskManagement(TrailingStopRiskManagementModel(0.04))
        self.SetWarmUp(timedelta(days=300))
    
    def beforeTheOpen(self):
        return
        
    def tradeStart(self):
        self.trade = True

    def tradeEnd(self):
        self.trade = False
        
    def buySignals(self):
        if self.trade == False:
            return
        
        # Return if benchmark is below SMA
        for symbolmark, symbolmarkData in self.marketDataBySymbol.items():
            if (self.Securities[symbolmark].Close <= symbolmarkData.sma200.Current.Value):
                return

        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not self.Portfolio[symbol].Invested and (self.Securities[symbol].Close > symbolData.sma200.Current.Value) and (symbolData.sma10.Current.Value > symbolData.sma20.Current.Value):
                self.stopMarketTicket = self.StopLimitOrder(symbol, 1, 0.9*symbolData.sma20.Current.Value, symbolData.sma20.Current.Value)
        
                    
            
    def sellSignals(self):
        if self.trade == False:
            return
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if self.Portfolio[symbol].Invested and (self.Securities[symbol].Close < symbolData.sma20.Current.Value):
                self.Liquidate(symbol, "Sell Signal")
        
class symbolMarkData:
    def __init__(self, symbol, sma50, sma200):
        self.Symbol = symbol
        self.sma50 = sma50
        self.sma200 = sma200

class SymbolData:
    def __init__(self, symbol, sma10, sma20, sma50, sma200, sma250, sma7, ema20, ema50, rsi, wilr, wilr_fast, atr, sma20_4hr, sma250_4hr):
        self.Symbol = symbol
        self.sma10 = sma10
        self.sma20 = sma20
        self.sma50 = sma50
        self.sma200 = sma200
        self.sma250 = sma250
        self.sma7 = sma7
        self.ema20 = ema20
        self.ema50 = ema50
        self.rsi = rsi
        self.wilr = wilr
        self.wilr_fast = wilr_fast
        self.atr = atr
        #self.emaConsolidate = emaConsolidate
        self.sma20_4hr = sma20_4hr
        self.sma250_4hr = sma250_4hr