from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel

class TrailingStopRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the maximum possible loss
    measured from the highest unrealized profit'''
    def __init__(self, maximumDrawdownPercent = 0.05):
        '''Initializes a new instance of the TrailingStopRiskManagementModel class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio compared with the highest unrealized profit, defaults to 5% drawdown'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.trailingHighs = dict()
        self.lastDay = -1
        
        self.percentGain = 0.005

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
            
        if algorithm.Time.day == self.lastDay:
            return []
        self.lastDay = algorithm.Time.day
        
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value
            percentChange = algorithm.Securities[symbol].Holdings.UnrealizedProfitPercent / 0.005

            # Add newly invested securities
            if symbol not in self.trailingHighs:
                self.trailingHighs[symbol] = security.Close    # Set to average holding cost
                continue
            
            # Remove if not invested
            if not security.Invested and symbol in self.trailingHighs:
                try:
                    self.trailingHighs.pop(symbol, None)
                except:
                    continue
                continue
            
            if percentChange.is_integer() and percentChange > 0:
                self.trailingHighs[symbol] = security.Close
            

            # Check for new highs and update - set to tradebar high
            # if self.trailingHighs[symbol] < security.High:
            #     self.trailingHighs[symbol] = security.High
            #     continue

            # Check for securities past the drawdown limit
            securityHigh = self.trailingHighs[symbol]
            if securityHigh == 0:
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                continue
            
            drawdown = (security.Low / securityHigh) - 1

            if drawdown < self.maximumDrawdownPercent:
                # liquidate
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))

        return riskAdjustedTargets

from itertools import groupby

class NadionResistanceShield(QCAlgorithm):



#class DataConsolidationAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        #self.SetEndDate(2019, 1, 1)
        self.SetCash(25000)  # Set Strategy Cash
        self.manually_selected = [self.AddEquity(symbol).Symbol for symbol in ["QQQ"]]
        self.tickers = self.ActiveSecurities
        # self.assets = filteredSymbols + self.manually_selected
        
        self.symbolDataBySymbol = {}
        self.MarketCaps = ["SPY", "TLT", "GLD", "VNQ"]# "QQQ"]#,"MDY","IWM"]
        self.marketDataBySymbol = {}
        
        self.volatilityDataBySymbol = {}
        self.vix = ["VIX"]
        
        self.trade = True
        self.atr=[]
        
        self.spy = "SPY"
        self.iwm = "IWM"
        self.mdy = "MDY"
        self.qqq = "QQQ"
        self.vix = "VIX"
        # Before the open
        
        # Trailing distance in $
        self.trail_dist = 10  
        
        # Declare an attribute that we shall use for storing our
        # stop loss ticket. 
        self.sl_order = None
        # Declare an attribute that we will use to store the last trail level
        # used. We will use this to decide whether to move the stop
        self.last_trail_level = None
   
        
          # Stock Selector
        self.AddUniverse(self.Coarse, self.Fine)
        self.activelyTrading = []
        self.weight = 0
        self.numberOfSymbolsCoarse = 500
        self.exposureToSector = 2
        self.lastMonth = -1
        
        
        for symbolmark in self.MarketCaps:
            symbol = self.AddEquity(symbolmark, Resolution.Hour).Symbol
            
            sma50 = self.SMA(symbol, 50, Resolution.Daily, Field.Close)
            sma200 = self.SMA(symbol, 200, Resolution.Daily, Field.Close)
            rsi = self.RSI(symbol, 14, Resolution.Daily)
            
            self.marketDataBySymbol[symbol] = symbolMarkData(symbol, sma50, sma200, rsi)
        
        for symbolvol in self.vix:
            symbol = self.AddEquity(symbolmark, Resolution.Hour).Symbol
            rsi = self.RSI(symbol, 14, Resolution.Daily)
            wilr = self.WILR(symbol, 14, Resolution.Daily)
            
            self.volatilityDataBySymbol[symbol] = symbolvolData(symbol, rsi, wilr)
        
        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'''
            
            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)
            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, 20, Resolution.Daily)
            self.atr.append(self.ATR(symbol, 7, Resolution.Daily))
        
            self.high = self.MAX(symbol, 5, Resolution.Daily, Field.High)
            self.midhigh = self.MAX(symbol, 3, Resolution.Daily, Field.High)

            self.low = self.MIN(symbol, 5, Resolution.Daily, Field.Low)
            self.stoplow = self.MIN(symbol, 20, 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)
            
            symbolData = SymbolData(symbol, ema10, sma20, sma200, sma7, sma50, ema20, ema50, rsi, wilr, wilr_fast, atr, smaConsolidate)
            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))
        
        

        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.AfterMarketOpen("SPY", 30), self.buySignals)         
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.AfterMarketOpen("SPY", 30), self.sellSignals)
        
        
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.BeforeMarketClose("SPY", 10), self.buySignals)
                
        self.Schedule.On(self.DateRules.EveryDay("SPY"),        
                self.TimeRules.BeforeMarketClose("SPY", 10), self.sellSignals)
                 
      

        #self.AddRiskManagement(TrailingStopRiskManagementModel(0.04))
        self.SetWarmUp(timedelta(days=180))
        
    def Coarse(self, coarse):
        if self.Time.month == self.lastMonth:
            return Universe.Unchanged
        self.lastMonth = self.Time.month
        
        allCoarse = [x for x in coarse if x.HasFundamentalData and x.Price > 1 and x.Volume > 1]
        finalCoarse = sorted(allCoarse, key = lambda x: x.DollarVolume, reverse = True)
        
        return [x.Symbol for x in finalCoarse][:self.numberOfSymbolsCoarse]
        return self.tickers
    
    def Fine(self, fine):
        
        filteredSymbols = []
        sortedBySector = [x for x in fine]
        for code, g in groupby(sortedBySector, lambda x: x.AssetClassification.MorningstarSectorCode):
            for x in sorted(g, key = lambda x: x.ValuationRatios.PERatio, reverse = False)[:self.exposureToSector]:
                filteredSymbols.append(x.Symbol)
        
        #self.tickers = filteredSymbols + self.manually_selected
    
        return filteredSymbols[:5]
        self.Log("tickers: {0}".format(self.tickers))
    
    
    def OnSecuritiesChanged(self, changes):
        return self.tickers
    
    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):
        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)))
        #self.Log(OrderEvent.FillPrice - symbolData.atr.Current.Value))

    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.rsi.Current.Value > 50):
                return
        
        for symbolvol, symbolvolData in self.volatilityDataBySymbol.items():
            if (self.Securities[symbolvol].Close >  symbolvolData.wilr.Current.Value < -20):
                return
            
        for symbol, symbolData in self.symbolDataBySymbol.items():
           if not self.Portfolio[symbol].Invested:#  and (self.Securities[symbol].Close < self.low.Current.Value):# and (self.midhigh.Current.Value < self.high.Current.Value):
                self.SetHoldings(symbol, .2, 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 > self.high.Current.Value):
                self.Liquidate(symbol, "Sell Signal")
          
                # Update our trailing stop loss as necessary
    def stopLoss(self):
        if  self.trade == False:
            return
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if self.Portfolio[symbol].Invested and (self.Securities[symbol].Close < (symbolData.atr.Current.Value *2)):
                self.Liquidate(symbol, "Sell 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, rsi):
        self.Symbol = symbol
        self.sma50 = sma50
        self.sma200 = sma200
        self.rsi = rsi

class symbolvolData:
    def __init__(self, symbol, rsi, wilr):
        self.Symbol = symbol
        self.rsi = rsi
        self.wilr = wilr


class SymbolData:
    def __init__(self, symbol, ema10, sma20, sma50, sma200, sma7, ema20, ema50, rsi, wilr, wilr_fast, atr, smaConsolidate):
        self.Symbol = symbol
        self.ema10 = ema10
        self.sma20 = sma20
        self.sma50 = sma50
        self.sma200 = sma200
        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.smaConsolidate = smaConsolidate