#region imports
from AlgorithmImports import *
#endregion
import numpy as np
from datetime import datetime
import talib

class NadionResistanceShield(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)  # Set Start Date
        #self.SetEndDate(2021, 1, 1)
        self.SetCash(25000)  # Set Strategy Cash
        self.tickers =     ['SOXX','CIBR','ROBO','CLOU','BLOK',                 #Tech
                                'IHI', 'VHT', 'IHE', 'EDOC',                        #Healthcare
                                'VDE', 'OIH', 'ICLN', 'XOP','TAN',                  #Energy
                                'XLF', 'IAT', 'FINX', 'IAI',                        #Financial
                                'GDX', 'LIT', 'COPX', 'VAW', 'URNM',"WOOD",         #Materials
                                'XLY', 'XLC','XRT', 'ESPO', 'CARZ',                 #Consumer Discretionary
                                'VIS', 'PHO', 'ITA', 'IYT', 'JETS',                 #Industrials
                                'XLU', 'PAVE', 'RYU', 'GIL'                         #Utilities
                                'VNQ', 'XLRE'                                       #RealEstate        
                                'VIXM', 'VXX', 'UVXY','BITS'                                       #Volatility
                                'QQQ', 'IWM', 'MDY', 'IEMG', 'DBC', 'IWF', 'IWD' 
         ]
        #self.ROC = {}
        self.num_fine = 5
        self.symbolDataBySymbol = {}
        self.marketDataBySymbol = {}
        self.vwaps = {}
        self.stds = {}
        self.Marketvwaps = {}
        self.Marketstds = {}

        self.trade = True
        self.atr=[]


        self.MarketCaps = ["SPY","QQQ"]#, "MDY","IWM"]
        self.spy = "SPY"
        self.iwm = "IWM"
        self.mdy = "MDY"
        self.qqq = "QQQ"
        self.momentum_periods = [30, 90, 180]
        self.consolidation_periods = [10, 20, 50]
        # trade only once per day
        self.next_trade_time = self.Time

        #self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))

   
        
        
        for symbolmark in self.MarketCaps:
            symbol = self.AddEquity(symbolmark, 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 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'''
            close = self.MIN(symbol, 1, Resolution.Daily, Field.Low)
            
            
            ema10 = self.EMA(symbol, 10, Resolution.Hour, 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)
            self.sma = self.SMA(symbol, 20, Resolution.Hour, Field.Close)
            sma50 = self.SMA(symbol, 50, Resolution.Daily, Field.Close)
            sma100 = self.SMA(symbol, 100, 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)
            
            
            self.vwaps[symbol] = self.VWAP(symbol, 14, Resolution.Daily)
            self.stds[symbol] = self.STD(symbol, 14, Resolution.Daily)  
            
            
            atr = self.ATR(symbol, 7, Resolution.Daily)
            self.atr.append(self.ATR(symbol, 7, Resolution.Daily))
        
            self.high = self.MAX(symbol, 7, Resolution.Daily, Field.High)
            self.low = self.MIN(symbol, 5, Resolution.Daily, Field.Low)
            
            self.sma.Updated += self.OnSMA
            
            '''Consolidator method'''
            smaConsolidate = ExponentialMovingAverage(20, 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, smaConsolidate, fourHourConsolidator)

            self.vwaps[symbol] = self.VWAP(symbol, 14, Resolution.Daily)
            self.stds[symbol] = self.STD(symbol, 14, Resolution.Daily)
            self.Marketvwaps[symbol] = self.VWAP(symbol, 14, Resolution.Daily)
            self.Marketstds[symbol] = self.STD(symbol, 14, Resolution.Daily)             

            
            symbolData = SymbolData(symbol, ema10, sma20, sma100, sma200, sma7, sma50, ema20, ema50, rsi, wilr, wilr_fast, atr, smaConsolidate, close)
            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=6.5)),
                 self.buySignals)
                 #self.TimeRules.EveryDay,
                 #self.buySignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=3.2)),
                 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.SetWarmUp(timedelta(days=90))
    
    def beforeTheOpen(self):
        self.Log("SPY: {0}".format(self.spy.Close))
        #for i in range(len(self.tickers)):
        #    self.Log("ATR: {0}".format(self.atr[i].Current.Value))
            
    
            
    
    def OnData(self, data):
        if self.IsWarmingUp:# or not all(std.IsReady for std in self.stds.values()):
            return

        for symbol in self.tickers:
            price = self.Securities[symbol].Price
            vwap = self.vwaps[symbol].Current.Value
            std = self.stds[symbol].Current.Value
            ub1 = vwap + std 
            ub2 = vwap + 2*std
            lb1 = vwap - std 
            lb2 = vwap - 2*std 

            self.Plot(symbol, 'price', price)
            self.Plot(symbol, 'vwap', vwap)
            self.Plot(symbol, 'ub2', ub2)
            self.Plot(symbol, 'ub1', ub1)
            self.Plot(symbol, 'lb1', lb1)
            self.Plot(symbol, 'lb2', lb2)
            
            #Take Profit
            if price > ub2:
                self.Liquidate(symbol)
                return
            

    def tradeStart(self):
        self.trade = True

    def tradeEnd(self):
        self.trade = False
    
    def OnOrderEvent(self, OrderEvent):
        '''Event when the order is filled. Debug log the order fill. :OrderEvent:'''

        if OrderEvent.FillQuantity == 0:
            return

        # Get the filled order
        Order = self.Transactions.GetOrderById(OrderEvent.OrderId)
        
        # Log the filled order details
        self.Log("ORDER NOTIFICATION >> {} >> Status: {} Symbol: {}. Quantity: "
                    "{}. Direction: {}. Fill Price {}".format(str(Order.Tag),
                                                   str(OrderEvent.Status),
                                                   str(OrderEvent.Symbol),
                                                   str(OrderEvent.FillQuantity),
                                                   str(OrderEvent.Direction),
                                                   str(OrderEvent.FillPrice)))
    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 < self.low.Current.Value) and (symbolData.sma100.Current.Value > symbolData.sma200.Current.Value):#  and (self.Securities[symbol].Close > symbolData.sma200.Current.Value):
                self.SetHoldings(symbol, .1, False, "Buy Signal")
                
           elif not self.Portfolio[symbol].Invested  and (self.Securities[symbol].Close < symbolData.sma20.Current.Value) and (symbolData.ema20.Current.Value > symbolData.ema50.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 (symbolData.rsi.Current.Value > 90):
                self.StopMarketOrder(symbol, .05, self.Securities[symbol].Close, "Some Profit Signal")
            #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 > self.high.Current.Value):
                self.Liquidate(symbol, "Sell New High Signal")
    
        
    def OnSMA(self, sender, updated):
        if self.sma.IsReady:
            #self.Debug(f"SMA Updated on {self.Time} with value: {self.sma.Current.Value}")
            return

    


class symbolMarkData:
    def __init__(self, symbol, sma50, sma200):
        self.Symbol = symbol
        self.sma50 = sma50
        self.sma200 = sma200

class SymbolData:
    def __init__(self, symbol, ema10, sma20, sma50, sma100, sma200, sma7, ema20, ema50, rsi, wilr, wilr_fast, atr, smaConsolidate, close):
        self.Symbol = symbol
        self.ema10 = ema10
        self.sma20 = sma20
        self.sma50 = sma50
        self.sma200 = sma200
        self.sma100 = sma100
        self.sma7 = sma7
        self.ema20 = ema20
        self.ema50 = ema50
        self.close = close 
        self.rsi = rsi
        self.wilr = wilr
        self.wilr_fast = wilr_fast
        self.atr = atr
        #self.emaConsolidate = emaConsolidate
        self.smaConsolidate = smaConsolidate