# 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):



#class DataConsolidationAlgorithm(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 =  [ "SPY"] 
        self.symbolDataBySymbol = {}
        self.trade = True
        self.atr=[]
        
        self.MarketCaps = ["IWM", "MDY", "SPY", "QQQ"]
        self.spy = ["SPY"]
        self.iwm = ["IWM"]
        self.mdy = ["MDY"]
        self.qqq = ["QQQ"]
        
    
        # Before the open
            
    
        
        
        
        for symbol in self.tickers:
            self.AddEquity(symbol, Resolution.Hour)
            
            '''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):
            for i in range(len(self.tickers)):
                self.Log("ATR: {0}".format(self.atr[i].Current.Value))
        
    def tradeStart(self):
        self.trade = True

    def tradeEnd(self):
        self.trade = False
        
    def buySignals(self):
        if self.trade == False:
            return
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
       #     if not (self.Securities[symbolmark].Close > symbolmarkData.sma200.Current.Value):
       #         elif symbol, symbolData in self.symbolDataBySymbol.items():
                    if not self.Portfolio[symbol].Invested and (self.Securities[self.spy].Close > symbolData.sma200.Current.Value) and (self.Securities[symbol].Close > symbolData.sma200.Current.Value) and (symbolData.sma10.Current.Value > symbolData.sma20.Current.Value):
                        self.stopMarketTicket = self.StopLimitOrder(symbol, self.CalculateOrderQuantity(symbol, 1), 0.9*symbolData.sma20.Current.Value, symbolData.sma20.Current.Value)
        
      #  for symbol, symbolData in self.symbolDataBySymbol.items():
           # if not self.Portfolio[symbol].Invested and (self.Securities[symbol].Close < symbolData.sma200.Current.Value) and (symbolData.rsi.Current.Value < 30) and (self.Securities[symbol].Close > symbolData.sma7.Current.Value):
           #    self.SetHoldings(symbol, .1, False, "Buy Signal")
           # elif not self.Portfolio[symbol].Invested  and (self.Securities[symbol].Close < symbolData.sma200.Current.Value) and (symbolData.rsi.Current.Value < 10):
           #     self.SetHoldings(symbol, .1, False, "Buy Signal")
      #     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.SetHoldings(symbol, 1, False, "Buy Signal")
     #           self.stopMarketTicket = self.StopLimitOrder(symbol, self.CalculateOrderQuantity(symbol, 1), 0.9*symbolData.sma20.Current.Value, symbolData.sma20.Current.Value)
           #elif not self.Portfolio[symbol].Invested  and (self.Securities[symbol].Close > symbolData.sma20.Current.Value) and (symbolData.sma7.Current.Value > symbolData.atr.Current.Value):
           #     self.SetHoldings(symbol, .1, False, "Buy Signal")
        
                    
            
    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.sma200.Current.Value) and (symbolData.sma20.Current.Value < symbolData.atr.Current.Value): 
               # self.Liquidate(symbol, "Sell Signal")
            if self.Portfolio[symbol].Invested and (self.Securities[symbol].Close < symbolData.sma20.Current.Value):
              #  self.stopLossTicket = self.StopLimitOrder(symbol, .1, self.Securities[symbol].Close  * 0.99)
                self.Liquidate(symbol, "Sell Signal")
           # elif self.Portfolio[symbol].Invested and (self.Securities[symbol].Close < symbolData.sma20.Current.Value) and (self.Securities[symbol].Close < symbolData.sma7.Current.Value):
            #    self.stopLossTicket = self.StopMarketOrder(symbol, .1, self.Securities[symbol].Close  * 0.985)
            #    self.Liquidate(symbol, "Sell Signal")
                # Update our trailing stop loss as necessary
        

    def OnData(self, slice):
        pass

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