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

from System import *
from System.Drawing import Color
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *

import numpy as np
import decimal as d
from datetime import timedelta, datetime

class ETFMomentumPython(QCAlgorithm):

    # private String
    #            0      1      2    3      4     5
    vSymbol = ["QQQ", "SPY", None, None, None, "TLT"]
    vArrNone = [None,None,None,None,None,None]
    vMaxRange = (6)
    
    # private MomentumPercent 
    Momentum = [None,None,None,None,None,None]

    # public ExponentialMovingAverage 
    emaFast = [None,None,None,None,None,None]
    emaSlow  = [None,None,None,None,None,None]

    # public AverageTrueRange 
    atr  = [None,None,None,None,None,None]

    # private int 
    momPer = (13)
    emaSlowPer = (195)
    emaFastPer = (55)
    atrPer = (8)

    def Initialize(self):
        self.SetStartDate(2006, 1, 1)  #Set Start Date
        self.SetCash(100000)       #Set Strategy Cash
        
        # Loop from 0 to 4
        for x in range (6):
            if self.vSymbol[x] is not None:
                self.AddEquity(self.vSymbol[x], Resolution.Daily)
                self.Momentum[x]  = self.MOMP(self.vSymbol[x], self.momPer, Resolution.Daily)
                self.emaFast[x] = self.EMA(self.vSymbol[x], self.emaFastPer, Resolution.Daily)
                self.emaSlow[x] = self.EMA(self.vSymbol[x], self.emaSlowPer, Resolution.Daily)
                self.atr[x] = self.ATR(self.vSymbol[x], self.atrPer, MovingAverageType.Exponential)

        self.SetBenchmark("SPY")
        self.SetWarmUp(self.emaSlowPer)
        
        Plot  = self.vArrNone
        # Loop from 0 to 4
        for x in range (6):
            if self.vSymbol[x] is not None:
                Plot[x] = Chart(self.vSymbol[x])
                Plot[x].AddSeries(Series("Buy", SeriesType.Scatter, "1", Color.Green, ScatterMarkerSymbol.Triangle))
                Plot[x].AddSeries(Series("Sell", SeriesType.Scatter, "1", Color.Red, ScatterMarkerSymbol.TriangleDown))
                Plot[x].AddSeries(Series("Price", SeriesType.Line, 0))
                Plot[x].AddSeries(Series("EMASlow", SeriesType.Line, 0))
                Plot[x].AddSeries(Series("EMAFast", SeriesType.Line, 0))
                Plot[x].AddSeries(Series("ATR", SeriesType.Line, 0))
                Plot[x].AddSeries(Series("MOMP", SeriesType.Line, 0))
                self.AddChart(Plot[x])

    def OnData(self, data):
        
        if (self.IsWarmingUp):
            return
                
        vemaSlow = [None,None,None,None,None,None]
        vemaFast = [None,None,None,None,None,None]
        vMom     = [None,None,None,None,None,None]
        vAtr     = [None,None,None,None,None,None]
        vScore   = [(0),(0),(0),(0),(0),(0)]
        vMaxScore   = None
        vMaxSymbol  = (self.vMaxRange-(1))
        vSetSymbol = None
        
        # auto plot our indicators on our symbol chart
        for x in range (self.vMaxRange):
            if self.vSymbol[x] is not None:
                vemaSlow[x] = self.emaSlow[x].Current.Value
                vemaFast[x] = self.emaFast[x].Current.Value
                vMom[x] = self.Momentum[x].Current.Value
                vAtr[x] = self.atr[x].Current.Value

        for x in range (self.vMaxRange):
            if self.vSymbol[x] is not None:
                if (data.ContainsKey(self.vSymbol[x])):
                    self.Plot(self.vSymbol[x], "EMASlow", vemaSlow[x])
                    self.Plot(self.vSymbol[x], "EMAFast", vemaFast[x])
                    #self.Plot(self.vSymbol[x], "ATR", vAtr[x])
                    self.Plot(self.vSymbol[x], "MOMP", vMom[x])
                    try:
                        self.Plot(self.vSymbol[x], "Price", data.Bars[self.vSymbol[x]].Close)
                    except:
                        pass

        for x in range (self.vMaxRange):
            if self.vSymbol[x] is not None:
                if vAtr[x] != 0:
                    vScore[x] = vMom[x]/vAtr[x]        

        for x in range (self.vMaxRange):
            if self.vSymbol[x] is not None:
                if vMaxScore is None:
                    if vScore[x] != 0:
                        vMaxScore = vScore[x]
                        vMaxSymbol = x
                elif vScore[x] > vMaxScore and vScore[x] != 0:
                    vMaxScore = vScore[x]
                    vMaxSymbol = x

        self.Log("{0} {1} {2} = {3}".format(vScore[0], vScore[1], vScore[5], self.vSymbol[vMaxSymbol]))
                    
        for x in range (self.vMaxRange):
            if self.vSymbol[x] is not None:
                if x != vMaxSymbol:
                    if (self.Portfolio[self.vSymbol[x]].Quantity > (0)):
                        self.Liquidate(self.vSymbol[x])
                        self.Plot(self.vSymbol[x], "Sell", -1)

        if (not self.Portfolio.Invested):
            vSetSymbol = None
            try:
                if (data.Bars[self.vSymbol[vMaxSymbol]].Close > vemaFast[vMaxSymbol] or data.Bars[self.vSymbol[vMaxSymbol]].Close > vemaSlow[vMaxSymbol]):
                    vSetSymbol = self.vSymbol[vMaxSymbol]
                    self.SetHoldings(self.vSymbol[vMaxSymbol], (0.99))
                    self.Plot(self.vSymbol[vMaxSymbol], "Buy", 1)

            except:
                vSetSymbol = self.vSymbol[vMaxSymbol]
                self.SetHoldings(self.vSymbol[vMaxSymbol], (0.99))
                self.Plot(self.vSymbol[vMaxSymbol], "Buy", 1)

        else:
            try:
                for x in range (self.vMaxRange):
                    if self.vSymbol[x] is not None:
                        if (self.Portfolio[self.vSymbol[x]].Quantity > (0)):
                            if (data.Bars[self.vSymbol[x]].Close < vemaFast[x] and data.Bars[self.vSymbol[x]].Close < vemaSlow[x]):
                                self.Liquidate(self.vSymbol[x])
                                self.Plot(self.vSymbol[x], "Sell", -1)
            except:
                pass

    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        #self.Log("OE: {0}: {1}: {2}".format(self.Time, order.Type, orderEvent))