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 *
import numpy as np
class VerticalCalibratedRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 3, 1)  # Set Start Date
        self.SetEndDate(2019, 4, 1)
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 5
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        self.smav = {}
        self.rsi = {}
        self.adx = {}
        self.atr = {}
        self.sma = {}
        self.selectedSymbols = []
        self._changes = None
        
    def OnData(self, data):
        pass

    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        
        filtered = [ x for x in coarse if x.HasFundamentalData and x.Price > 1 and x.DollarVolume > 2500000 ]
        sortedByDollarVolume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        
        # return the symbol objects of the top entries from our sorted collection
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]

    
    # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
    def FineSelectionFunction(self, fine):
        # sort descending by P/E ratio
        sortedByPeRatio = sorted(fine, key=lambda x: x.ValuationRatios.PERatio, reverse=True)

        ## Retrieve 20 days of historical data for each symbol
        symbols = [x.Symbol for x in sortedByPeRatio]
        history = self.History(symbols, 150, Resolution.Daily)
        ## Iterate through symbols
        for symbol in symbols:
            ## Find hsitory for specific symbol
            if str(symbol) in history.index.get_level_values(0):
                symbolVolumeHistory = history.loc[str(symbol)]
                if symbolVolumeHistory.empty:
                    self.Log("EMPTY dataframe!")
    
                ## Create SMA for symbol and register it with algorithm
                else:
                    symbolSMAv = SimpleMovingAverage(50)
                    symbolRSI = RelativeStrengthIndex(3)
                    symbolADX = AverageDirectionalIndex(7)
                    symbolATR = AverageTrueRange(10)
                    symbolSMA = SimpleMovingAverage(150)
    
                ## Iterate through historical data
                for tuple in symbolVolumeHistory.itertuples():
                    bar = TradeBar(tuple.Index, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)
                    ## Update SMA with data time and volume
                    symbolSMAv.Update(tuple.Index, tuple.volume)
                    #self.Debug(f'Updating {symbol.Value} SMA...')
                    symbolRSI.Update(tuple.Index, tuple.close)
                    symbolADX.Update(bar)
                    symbolATR.Update(bar)
                    symbolSMA.Update(tuple.Index, tuple.close)
                    
                    self.Log(f'{symbol.Value} ADX: {symbolADX.Current.Value}')
                    self.Log(f'{symbol.Value} ATR: {symbolATR.Current.Value}')
                    
                ## Add SMA to dictionary so you can access it later
                self.smav[symbol] = symbolSMAv
                self.rsi[symbol] = symbolRSI
                self.adx[symbol] = symbolADX
                self.atr[symbol] = symbolATR
                self.sma[symbol] = symbolSMA
                #self.Debug(f' {symbol.Value} Price ' + str(symbolVolumeHistory['close'][-1]) + 'SMA' + str(self.sma[symbol]))
            else:
                symbols.pop(symbols.index(symbol))

        ## Perform SMA filtering conditions and select/return the symbols you want to add to your universe
        ## Fine Selection will use only these ones
        
        self.selectedSymbols = [ x for x in symbols if self.smav[x].Current.Value > 500000 and self.rsi[x].Current.Value < 30 and history.loc[str(x)]['close'][-1] > self.sma[x].Current.Value  ]
        
        return self.selectedSymbols
        
    def OnData(self, data):
        # if we have no changes, do nothing
        if self._changes is None: return

        # liquidate removed securities
        for security in self._changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)

        # we want 20% allocation in each security in our universe
        for security in self._changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.1)

        self._changes = None


    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes