#region imports
from AlgorithmImports import *
#endregion
from Alphas.MacdAlphaModel import MacdAlphaModel

class WellDressedYellowGreenFish(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        self.SetCash(5000)  # Set Strategy Cash
        #self.AddAlpha(MacdAlphaModel(12, 26, 9, MovingAverageType.Simple, Resolution.Daily))
        self.SetBenchmark("SPY")
        self.AddUniverse(self.CoarseSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.averages = { }
        self.SetWarmUp(400)
    
    def CoarseSelectionFunction(self, universe):  
        selected = []
        universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)  
        universe = [c for c in universe if c.Price > 100][:100]

        for coarse in universe:  
            symbol = coarse.Symbol
            
            if symbol not in self.averages:
                # 1. Call history to get an array of 200 days of history data
                history = self.History(symbol, 200, Resolution.Daily)
                
                #2. Adjust SelectionData to pass in the history result
                self.averages[symbol] = SelectionData(history) 

            self.averages[symbol].update(self.Time, coarse.AdjustedPrice)
            
            if  self.averages[symbol].is_ready() and self.averages[symbol].fast > self.averages[symbol].mid and self.averages[symbol].mid >  self.averages[symbol].slow:
                selected.append(symbol)
        
        return selected[:2]
        
    def OnData(self, data):
        for symbol, symbol_data in self.averages.items():
            tolerance = 0.0015
            if not data.ContainsKey(symbol) or not data.Bars.ContainsKey(symbol):
                continue
            holdings = self.Portfolio[symbol].Quantity if self.Portfolio.ContainsKey(symbol) else 0
            signalDeltaPercent = (symbol_data.macd.Current.Value - symbol_data.macd.Signal.Current.Value)/symbol_data.macd.Fast.Current.Value
            if holdings <= 0 and signalDeltaPercent > tolerance:  # 0.01%
                # longterm says buy as well
                self.SetHoldings(symbol, 0.5)
            elif holdings >= 0 and signalDeltaPercent < -tolerance:
                self.Liquidate(symbol)
                

class SelectionData():
    #3. Update the constructor to accept a history array
    def __init__(self, history):
        self.slow = ExponentialMovingAverage(200)
        self.mid = ExponentialMovingAverage(13)
        self.fast = ExponentialMovingAverage(5)
        self.macd = MovingAverageConvergenceDivergence(50, 200, 9, MovingAverageType.Exponential)
        
        #4. Loop over the history data and update the indicators
        for bar in history.itertuples():
            self.fast.Update(bar.Index[1], bar.close)
            self.slow.Update(bar.Index[1], bar.close)
            self.mid.Update(bar.Index[1], bar.close)
            self.macd.Update(bar.Index[1], bar.close)
            
    def is_ready(self):
        return self.slow.IsReady and self.fast.IsReady and self.mid.IsReady
        
    def update(self, time, price):
        self.fast.Update(time, price)
        self.mid.Update(time, price)
        self.slow.Update(time, price)
        self.macd.Update(time, price)