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

from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
from datetime import datetime, timedelta

import datetime

class ConsolidatorHelp(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 12, 1)  # Set Start Date
        
        self.SetEndDate(datetime.datetime.now())
        
        self.SetCash(600000)  # Set Strategy Cash
        self.SetWarmup(20)
        self.SetBenchmark("SPY")
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.UniverseSettings.Leverage = 2
        self._changes = None
        self.MinData = {}
        self.dayMinData = {}
        
        self.Data = {} # dictionary to hold rolling window for each symbol
        self.window_hr = {}
        self.windowmin = {}
        
        # Set params for tracking stop losses
        # -----------------------------------------------
        self.activateStopCoef     = 1   # activate trail when price rises 'x * ATR' 
        self.initialStopCoef      = 3.5     # exit when price dips 'x * ATR'
        self.trailStopCoef        = 11       # exit when price dips 'x * ATR'(trailing)
        self.trailActivationPrice = 0
        self.trailActivationPriceSells = 0
        self.trailStopActivated   = False     
        self.trailStopActivatedSells   = False 
        self.trailingStopLoss     = 0 
        self.initialStopLoss      = 0
        self.trailingStopLossSells     = 0 
        self.initialStopLossSells      = 0
        
        self.ADXtresh             = 20  
        
        ##___________FILTERS_________________##
        
    def CoarseSelectionFunction(self, coarse ):
        
        
        top = [x for x in coarse if x.HasFundamentalData
                           and 500 > x.Price > 40
                           and  x.Volume > 500000]
        self.Debug("len of coarse = " + str(len(top)))
        return [x.Symbol for x in top]
        
    def FineSelectionFunction(self, fine ):
        
        top = [x for x in fine if  x.ValuationRatios.PERatio > 20 and
                           x.ValuationRatios.TrailingDividendYield < .02 and 
                           not x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.BasicMaterials and
                           not x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Energy]
        top = sorted(top, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        self.Debug("len of fine = " + str(len(top)))
        
        return  [x.Symbol for x in top][:5]
        
    def OnData(self, data):
        
         # Make sure indicators and rolling windows are ready
        if not all([symbol.IsReady for symbol in self.dayMinData.values()]):
                    return
         
        if not all([symbol.IsReady for symbol in self.MinData.values()]):
            return
            
            '''
        for symbol, symbolData in self.dayMinData.items():
            if not symbolData.IsReady:
                    return
        
        for symbol, symbolData in self.MinData.items():
            if not symbolData.IsReady:
                    return 
        
        for symbol in self.dayMinData.items():
            if not (self.ema.IsReady and self.emaWindow.IsReady and \
                self.rsi.IsReady and self.rsiWindow.IsReady and \
                self.BB1UpperWindow.IsReady and self.BB1LowerWindow.IsReady and self.window_hr.IsReady):
                    return
            
        self.Debug("Are WE GETTING DATA???")'''
     
        #---------------------------------------------------------------_--_#
        self.windowmin.Add(data[symbol])
        self.window_hr.Add(data[symbol])
        
        for symbol, value in self.dayMinData.items():
            
            if  data.ContainsKey(symbol) and \
                data[symbol] != None:
                self.Debug("AM I GETTING DATA???")
                self.Debug(str(symbol))    
                
        for symbol, value in self.MinData.items():
            if (value.windowmin[0].EndTime.hour < 15 or (value.windowmin[0].EndTime.hour == 15 and value.windowmin[0].EndTime.minute < 30)) and \
                data.ContainsKey(symbol) and not self.Portfolio[symbol].Invested and \
                value.windowmin[0].EndTime == self.Time and \
                data[symbol] != None:
        
                self.windowmin.Add(data[symbol])
                self.window_hr.Add(data[symbol])
                self.Debug("I am GETTING DATA")
                self.Debug(str(symbol))
                
        self.Debug("Windows ready one")
        
        currBar = self.window_hr[0].Close
        close_of_previous_bar = self.window_hr[1].Close
        self.Log(f"Latest `window` close for {symbol}: {close_of_previous_bar}")
            
        currMinCloseBar = self.windowmin[0].Close
        close_of_previous_Min_bar = self.windowmin[1].Close
        
        currema = self.emaWindow[0]                     # Current ema had index zero.
        pastema = self.emaWindow[self.emaWindow.Count-1]   # Oldest ema has index of window count minus 1.
        self.Log("ema:   {0} -> {1} ... {2} -> {3}".format(pastema.Time, pastema.Value, currema.Time, currema.Value))
        
        currrsi = self.rsiWindow[0]                     # Current rsi had index zero.
        pastrsi = self.rsiWindow[self.rsiWindow.Count-1]   # Oldest rsi has index of window count minus 1.
        
        currBBUpW = self.BB1UpperWindow[0] # Current bbup had index zero.
        pastBBUpW = self.BB1UpperWindow[self.BB1UpperWindow.Count-1] # Oldest bbup has index of window count minus 1.
        currBBlW = self.BB1LowerWindow[0] # Current bbdown had index zero.
        pastBBlW = self.BB1LowerWindow[self.BB1LowerWindow.Count-1] # Oldest bbdown has index of window count minus 1.
        
        self.diplus = self.ADX_hr.PositiveDirectionalIndex.Current.Value
        self.diminus = self.ADX_hr.NegativeDirectionalIndex.Current.Value
        
        if not (self.Portfolio.Invested and currema.Value > pastema.Value and self.diplus > self.ADXtresh): 
            if currrsi.Value < 30 and \
                currMinCloseBar > currBBUpW:
                    self.SetHoldings(symbol, 0.5)
                    self.SetInitialStopsBuys()
                    
        if not (self.Portfolio.Invested and currema.Value < pastema.Value and self.diminus > self.ADXtresh): 
            if currrsi.Value > 70 and \
                currMinCloseBar < currBBlW:
                    self.SetHoldings(symbol, -0.5)
                    self.SetInitialStopsSells()
                    
    # Add and remove stocks whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            # For 5 minute quotes
            if symbol not in self.MinData: 
                self.MinData[symbol] = SymbolData(self, symbol) 
                
            
            # For daily quotes
            if symbol not in self.dayMinData:
                self.dayMinData[symbol] = DaySymbolData(self, symbol)
                
        
        '''for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.MinData: #self.Data:
                symbolData = self.MinData.pop(symbol, None) 
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.fiveMinutesConsolidator)
            if symbol in self.DayMinData: #self.Data:
                symbolData = self.DayMinData.pop(symbol, None) 
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.HourConsolidator)'''
                
       ##__________5 min Quotes ______##
   
class SymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        self.windowmin = RollingWindow[QuoteBar](20)
        
        fiveMinutesConsolidator = QuoteBarConsolidator(timedelta(minutes=5))
        fiveMinutesConsolidator.DataConsolidated += self.fiveMinutesBarHandler 
        algorithm.SubscriptionManager.AddConsolidator(symbol, fiveMinutesConsolidator)# Register consolidator to get automatically updated with minute data
        
        ## Create nd Add Indicators to rolling window
            
        self.ema = ExponentialMovingAverage(symbol, 20)
        self.ema.Updated += self.emaUpdated
        algorithm.RegisterIndicator(symbol, self.ema, fiveMinutesConsolidator)
        self.emaWindow = RollingWindow[IndicatorDataPoint](21)
        
        self.rsi = RelativeStrengthIndex(symbol, 14)
        self.rsi.Updated += self.RsiUpdated
        algorithm.RegisterIndicator(symbol, self.rsi, fiveMinutesConsolidator)
        self.rsiWindow = RollingWindow[IndicatorDataPoint](21)
        
        self.macd3 = MovingAverageConvergenceDivergence(3, 6, MovingAverageType.Exponential)
        algorithm.RegisterIndicator(symbol, self.macd3, fiveMinutesConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(symbol, fiveMinutesConsolidator)
        
           # -------------BB bands_________________#
           
        self.BB1 = BollingerBands(20, 2)
        self.BB1UpperWindow=RollingWindow[float](20)
        self.BB1LowerWindow=RollingWindow[float](20)
        self.BB1.Updated+=self.BB1Updated
        algorithm.RegisterIndicator(symbol, self.BB1, fiveMinutesConsolidator)
           
    # Update Indicator Windows
    
    def emaUpdated(self, sender, updated):
        '''Adds updated values to rolling window'''
        if self.ema.IsReady:
            self.emaWindow.Add(updated)
        
    def RsiUpdated(self, sender, updated):
            if self.rsi.IsReady:
                self.rsiWindow.Add(updated)# Add updated indicator data to rolling window
    
    def BB1Updated(self, sender, updated):
        self.BB1UpperWindow.Add(sender.UpperBand.Current.Value)
        self.BB1LowerWindow.Add(sender.LowerBand.Current.Value)
        
    # Lock and load
    @property    
    def IsReady(self):
        if self.ema.IsReady and self.emaWindow.IsReady and \
            self.rsi.IsReady and self.rsiWindow.IsReady and \
            self.BB1UpperWindow.IsReady and self.BB1LowerWindow.IsReady:
                return
            
    # Make sure I can get stock values in Window.
    def fiveMinutesBarHandler(self, sender, bar):
        self.windowmin.Add(bar)   
        '''
        self.algorithm.Debug(f"Data Consolidatoed for {self.symbol} \
            at {bar.EndTime} with bar: {bar} and sma {self.sma} and trix {self.trix} and \
            kelt {self.kelt.LowerBand}")
        '''
    
     ##__________60 Mins Quotes ______## 
     
class DaySymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        ##########  INDICATORS  #################################
        self.window_hr = RollingWindow[QuoteBar](20)
        
        
        ########### CONSOLIDATORS and WINDOWS ###########################
        HourConsolidator = QuoteBarConsolidator(timedelta(hours=1)) #BarPeriod = TimeSpan.FromMinutes(10)
        HourConsolidator.DataConsolidated += self.HourBarHandler
        algorithm.SubscriptionManager.AddConsolidator(symbol, HourConsolidator)# Register consolidator to get automatically updated with minute data
        
        ##-------- OTHER INDICATORS ----------##
        self.ATR_init     = AverageTrueRange(symbol, 5)
        self.ATR_trail    = AverageTrueRange(symbol, 10)
        self.ATR_activate = AverageTrueRange(symbol, 15)
        self.ADX_hr         = AverageDirectionalIndex(symbol, 10)

    
    # Lock and load
    @property    
    def IsReady(self):
        return self.window_hr.IsReady
    
    # Make sure I can get stock values in Window.
    def HourBarHandler(self, sender, bar):
        self.window_hr.Add(bar)
        #self.algorithm.Debug(f"Data Consolidated for {self.symbol} \
            #at {bar.EndTime} with bar: {bar} ")
        
     # ========================================================================
    # Set initial stop and activation level. Called after new position opened.
    # ========================================================================
    def SetInitialStopsBuys(self):
        currBar = self.window_hr[0].Close
        self.atrValue_init        = self.ATR_init.Current.Value
        self.atrValue_activate    = self.ATR_activate.Current.Value
        self.trailStopActivated   = False
        self.initialStopLoss      = currBar - (self.atrValue_init  * self.initialStopCoef)
        self.trailActivationPrice = currBar + (self.atrValue_activate * self.activateStopCoef)
        #self.PlotCharts()         # Plot charts for debugging
    
    def SetInitialStopsSells(self):
        currBar = self.window_hr[0].Close
        self.atrValue_initsells        = self.ATR_init.Current.Value
        self.atrValue_activatesells    = self.ATR_activate.Current.Value
        self.trailStopActivatedSells   = False
        self.initialStopLossSells      = currBar + (self.atrValue_initsells  * self.initialStopCoef)
        self.trailActivationPriceSells = currBar - (self.atrValue_activatesells * self.activateStopCoef)