# https://quantpedia.com/Screener/Details/25
from System.Drawing import Color
from Alphas.EmaCrossAlphaModel import EmaCrossAlphaModel

class SmallCapInvestmentAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2020, 1, 8)
        self.SetEndDate(2021, 1, 10)
        self.SetCash(100000)

        self.lastMonth = -1
        
        self.emaLongestWindow = 200
        self.UniverseSettings.MinimumTimeInUniverse = self.emaLongestWindow
        self.SetWarmUp(self.emaLongestWindow)
        
        self.coarseCount = 100
        self.fineCount = 10

        self.indicators = {}

        self.UniverseSettings.Resolution = Resolution.Hour
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        stockPlot = Chart("Trade Plot")
        stockPlot.AddSeries(Series('Price', SeriesType.Line, '$', Color.Blue))
        stockPlot.AddSeries(Series(" 200EMAIndicator", SeriesType.Line, '$', Color.Yellow))
        
        stockPlot = Chart("Diamonds")
        stockPlot.AddSeries(Series('coarseSelect', SeriesType.Scatter, '$', Color.Gold, ScatterMarkerSymbol.Diamond))
        stockPlot.AddSeries(Series('fineSelect', SeriesType.Scatter, '$', Color.Purple, ScatterMarkerSymbol.Diamond))
        stockPlot.AddSeries(Series('Above220EMA', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Diamond))
        
        self.AddChart(stockPlot)
        
        coarseSelectAction = 0
        self.Plot("Diamonds", "coarseSelect", coarseSelectAction)
        fineSelectAction = 0
        self.Plot("Diamonds", "fineSelect", fineSelectAction)

    def CoarseSelectionFunction(self, coarse):
        ''' Drop stocks which have no fundamental data or have low price '''
        if self.lastMonth == self.Time.month:
            return Universe.Unchanged
        self.lastMonth = self.Time.month
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        self.symbols = [ x.Symbol for x in sortedByDollarVolume[:self.coarseCount] ]
        coarseSelectAction = 45
        self.Plot("Diamonds", "coarseSelect", coarseSelectAction)
        return self.symbols
 
 
    def FineSelectionFunction(self, fine):
        ''' Selects the stocks by lowest market cap '''
        sortedByPeRatio = sorted(fine, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        self.symbols = [ x.Symbol for x in sortedByPeRatio[:self.fineCount] ]
        fineSelectAction = 50
        self.Plot("Diamonds", "fineSelect", fineSelectAction)
        #return [x.Symbol for x in sorted_market_cap[:self.count]]
        return self.symbols


    def OnData(self, data):

        for symbol in self.symbols:
            if not data.ContainsKey(symbol):
                continue
            if data[symbol] is None:
                continue
            
            if symbol not in self.indicators:
                self.indicators[symbol] = SymbolData(symbol, self, self.emaLongestWindow)
            # update by bar
            #self.indicators[symbol].update_bar(data[symbol])
            #update by value
            self.indicators[symbol].update_value(self.Time, data[symbol].Price)
            
            if self.IsWarmingUp: continue
            
            self.Log(str(symbol) + " : " + str(self.indicators[symbol].get_longestEMA()))
            
            sym_price = data[symbol].Price
            
            #self.EMAreq = self.EMA(symbol, 220, Resolution.Hour);
            #self.PlotIndicator('Trade Plot', self.EMAreq);
            self.Plot('Trade Plot',symbol.Value + " " + 'Price', data.Bars[symbol].Close)
            
            longestEMA = self.indicators[symbol].get_longestEMA()
            self.Plot("Trade Plot", (symbol.Value + " " + "200EMAIndicator"), self.indicators[symbol].get_longestEMA());
            if longestEMA != 0.0:
                if longestEMA > sym_price:
                    self.SetHoldings(symbol, -0.1)
                    emaclear = 0
                    self.Plot("Diamonds", "Above220EMA", emaclear)
                else:
                    self.Liquidate(symbol)


    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if self.Portfolio[symbol].Invested:
                self.Liquidate(symbol, 'Removed from Universe')
                # clean up
                del self.indicators[security.Symbol]

class SymbolData(object):
    def __init__(self, symbol, context, window):
        self.symbol = symbol

        self.window    = window
        self.indicator = context.EMA(symbol, self.window)
        self.longestEMA = 0.0

    def update_bar(self, bar):
        self.indicator.Update(bar)
        
    def update_value(self, time, value):
        self.indicator.Update(time, value)
            
    def get_longestEMA(self):
        return self.indicator.Current.Value