# https://quantpedia.com/Screener/Details/100
from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Data import SubscriptionDataSource
from QuantConnect.Algorithm import *
from QuantConnect.Python import PythonData, PythonQuandl
from sklearn import datasets, linear_model
from datetime import date, timedelta, datetime
from collections import deque
import statsmodels.api as sm
import decimal
import numpy as np

class TradeWtiBrentSpreadAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2018, 12, 15)#datetime.now())
        self.SetCash(1000000)
        # import the custom data 
        '''self.AddData(WTI, "WTI", Resolution.Daily)
        self.AddData(BRENT, "BRENT", Resolution.Daily)'''
        self.es = "SCF/CME_ES1_ON"
        self.AddData(QuandlFuture, self.es, Resolution.Daily)
        self.nq = "SCF/CME_NQ1_ON"
        self.AddData(QuandlFuture, self.nq, Resolution.Daily)
        # Subscribe and set our expiry filter for the futures chain
        futureES = self.AddFuture(Futures.Indices.SP500EMini)
        futureES.SetFilter(timedelta(0), timedelta(90))
        futureNQ = self.AddFuture(Futures.Indices.NASDAQ100EMini)
        futureNQ.SetFilter(timedelta(0), timedelta(90))
        benchmark = self.AddEquity("SPY");
        self.SetBenchmark(benchmark.Symbol);
        self.frontES = None
        self.frontNQ = None
        
        #hist = self.History(["WTI", "BRENT"], 400, Resolution.Daily)["value"].unstack(level=0).dropna()
        #Create the EMA spread for the liquidate signal
        '''self.wtiema = self.EMA("WTI", 7)
        self.brentema = self.EMA("BRENT", 7)
        self.wtibrentema = IndicatorExtensions.Minus(self.wtiema, self.brentema)
        self.SpreadSMA = SimpleMovingAverage(13)#for WTI|BRENT'''
        
        self.esema = self.EMA(self.es, 8)
        self.nqema = self.EMA(self.nq, 8)
        self.esnqema = IndicatorExtensions.Minus(self.nqema, self.esema)
        self.esnqSMA = SimpleMovingAverage(13)
        
        # Add the spread plot and mark the long/short spread point
        spreadPlot = Chart("Spread Plot")
        #spreadPlot.AddSeries(Series("Spread", SeriesType.Line, 0))
        #spreadPlot.AddSeries(Series("Long Spread Trade", SeriesType.Scatter, 0))
        #spreadPlot.AddSeries(Series("Short Spread Trade", SeriesType.Scatter, 0))
        #spreadPlot.AddSeries(Series("Spread EMA", SeriesType.Line, 0))
        #spreadPlot.AddSeries(Series("Spread SMA", SeriesType.Line, 0))
        spreadPlot.AddSeries(Series("Spread2", SeriesType.Line, 0))
        spreadPlot.AddSeries(Series("ESNQ EMA", SeriesType.Line, 0))
        spreadPlot.AddSeries(Series("ESNQ SMA", SeriesType.Line, 0))
        self.AddChart(spreadPlot)
        self.SetWarmup(timedelta(100))
    def OnData(self, data):
        #self.Log("Warming Up")
        if self.IsWarmingUp: return #for some reason this was stuck warming up...i added the time delta bit, see what happens next..
        '''self.Log("checcking for data contains key wti brent")
        if (data.ContainsKey("WTI") and data.ContainsKey("BRENT")): 
            
            self.SpreadSMA.Update(self.Time, data["WTI"].Price - data["BRENT"].Price) 
        if not self.SpreadSMA.IsReady: return'''
        #self.Plot("Spread Plot", "Spread", data["WTI"].Price - data["BRENT"].Price)
        #self.Log("checking for data contains key ES NQ")
        if (data.ContainsKey("SCF/CME_NQ1_ON") and data.ContainsKey("SCF/CME_ES1_ON")):
            #self.Log("checking the nasdaq sma and plotting spread")
            self.esnqSMA.Update(self.Time, data["SCF/CME_NQ1_ON"].Price - data["SCF/CME_ES1_ON"].Price)
            self.Plot("Spread Plot", "Spread2", data["SCF/CME_NQ1_ON"].Price - data["SCF/CME_ES1_ON"].Price)
            
        #self.Log("checking if ES NQ SMA is ready")
        if not self.esnqSMA.IsReady: return
        #self.Plot("Spread Plot", "Spread EMA", self.wtibrentema.Current.Value)
        #self.Plot("Spread Plot", "Spread SMA", self.SpreadSMA.Current.Value)
        tolerance = decimal.Decimal(0.005);
        #pread = self.Securities["WTI"].Price - self.Securities["BRENT"].Price
        spread2 = self.Securities[self.nq].Price - self.Securities[self.es].Price
        '''
        if spread > self.SpreadSMA.Current.Value and not (self.Portfolio["WTI"].IsShort and self.Portfolio["BRENT"].IsLong):
            self.SetHoldings("WTI", -0.2)
            self.SetHoldings("BRENT", 0.2)
            #self.Plot("Spread Plot", "Long Spread Trade", data["WTI"].Price - data["BRENT"].Price)
            #self.Debug()
            
        elif spread < self.SpreadSMA.Current.Value and not (self.Portfolio["WTI"].IsLong and self.Portfolio["BRENT"].IsShort):
            self.SetHoldings("WTI", 0.2)
            self.SetHoldings("BRENT", -0.2)
            #self.Plot("Spread Plot", "Short Spread Trade", data["WTI"].Price - data["BRENT"].Price)
            
        if self.Portfolio["WTI"].IsShort and self.Portfolio["BRENT"].IsLong and spread < self.wtibrentema.Current.Value:
            self.Liquidate()
            #self.Debug("the ema of wti/brent is " + str(self.wtibrentema.Current.Value))
        
        if self.Portfolio["WTI"].IsLong and self.Portfolio["BRENT"].IsShort and spread > self.wtibrentema.Current.Value:
            self.Liquidate()'''
        #for kvp in data.Bars:
        #    self.Debug("---> OnData: {}, {}, {}" .format(self.Time, kvp.Key.Value, kvp.Value.Close));
        self.Plot("Spread Plot", "ESNQ SMA", self.esnqSMA.Current.Value)
        self.Plot("Spread Plot", "ESNQ EMA", self.esnqema.Current.Value)
        for chain in data.FutureChains:
            if chain.Key.Value == Futures.Indices.SP500EMini:
                if self.frontES is None :# Get contracts expiring no earlier than in 90 days
                    contracts = list(filter(lambda x: x.Expiry > self.Time + timedelta(10), chain.Value))
            # if there is any contract, trade the front contract
                    if len(contracts) == 0: continue
                    self.frontES = sorted(contracts, key = lambda x: x.Expiry, reverse=True)[0]
        
            if chain.Key.Value == Futures.Indices.NASDAQ100EMini:
                if self.frontNQ is None:
                    # Get contracts expiring no earlier than in 90 days
                    contracts = list(filter(lambda x: x.Expiry > self.Time + timedelta(10), chain.Value))
                # if there is any contract, trade the front contract
                    if len(contracts) == 0: continue
                    self.frontNQ = sorted(contracts, key = lambda x: x.Expiry, reverse=True)[0]
        #self.Debug("self.frontES is " + str(self.frontES))
        #self.Debug("self.frontNQ is " + str(self.frontNQ))
        if self.frontES and self.frontNQ:
            #self.Log(str(spread2) + " is spread2 and " + str(self.esnqSMA.Current.Value) + " is ESNQSimpleMovingAverage")
            if spread2 > self.esnqSMA.Current.Value * (1 + tolerance) and not (self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong):       
                #self.Log(f'frontNQ is short and frontES is long: {self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong}')  
                self.SetHoldings(self.frontNQ.Symbol , -0.2)
                self.SetHoldings(self.frontES.Symbol , 0.2)    
                #self.Log("Absolute HOldings cost is " + str(self.Portfolio.TotalAbsoluteHoldingsCost))
                '''if (data.ContainsKey("SCF/CME_NQ1_ON") and data.ContainsKey("SCF/CME_ES1_ON")):    
                    self.Plot("Spread Plot", "Long Spread Trade", data[self.nq].Price - data[self.es].Price)'''
            elif spread2 < self.esnqSMA.Current.Value * (1 - tolerance) and not (self.Portfolio[self.frontNQ.Symbol].IsLong and self.Portfolio[self.frontES.Symbol].IsShort):
                #self.Log(f'frontNQ is short and frontES is long: {self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong}')
                self.SetHoldings(self.frontNQ.Symbol , 0.2)
                self.SetHoldings(self.frontES.Symbol , -0.2)    
                #self.Log("Absolute HOldings cost is " + str(self.Portfolio.TotalAbsoluteHoldingsCost))
                '''if (data.ContainsKey("SCF/CME_NQ1_ON") and data.ContainsKey("SCF/CME_ES1_ON")):
                    self.Plot("Spread Plot", "Short Spread Trade", data[self.nq].Price - data[self.es].Price)'''
            if self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong and spread2 < self.esnqema.Current.Value:
                #self.Log(f'Is it time to Liquidate?: {self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong}')  
                self.Liquidate(self.frontES.Symbol)
                self.Liquidate(self.frontNQ.Symbol)
                #self.Log("Liquidated")
            if self.Portfolio[self.frontNQ.Symbol].IsLong and self.Portfolio[self.frontES.Symbol].IsShort and spread2 > self.esnqema.Current.Value:
                #self.Log(f'Is it time to Liquidate?: {self.Portfolio[self.frontNQ.Symbol].IsShort and self.Portfolio[self.frontES.Symbol].IsLong}')
                self.Liquidate(self.frontES.Symbol)
                self.Liquidate(self.frontNQ.Symbol)
                #self.Log("Liquidated")
        
    '''def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))'''
            
class QuandlFuture(PythonQuandl):
    '''Custom quandl data type for setting customized value column name. Value column is used for the primary trading calculations and charting.'''
    def __init__(self):
        # Define ValueColumnName: cannot be None, Empty or non-existant column name
        # If ValueColumnName is "Close", do not use PythonQuandl, use Quandl:
        # self.AddData[QuandlFuture](self.crude, Resolution.Daily)
        self.ValueColumnName = "Settle"             
class WTI(PythonData):
    "Class to import WTI Spot Price(Dollars per Barrel) data from Dropbox"
    
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("https://www.dropbox.com/s/jpie3z6j0stp97d/wti-crude-oil-prices-10-year-daily.csv?dl=1", SubscriptionTransportMedium.RemoteFile)

    def Reader(self, config, line, date, isLiveMode):
        if not (line.strip() and line[1].isdigit()): return None
        index = WTI()
        index.Symbol = config.Symbol
        try:
            # Example File Format: (Data starts from 08/11/2008)
            # date     value
            # 8/11/08    114.44
            data = line.split(',')
            index.Time = datetime.strptime(data[0], "%Y-%m-%d")
            index.Value = Decimal(data[1])
        except:
            return None
            
        return index

class BRENT(PythonData):
    "Class to import BRENT Spot Price(Dollars per Barrel) data from Dropbox"
    
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("https://www.dropbox.com/s/w380c4n7xjmdqxl/brent-crude-oil-prices-10-year-daily.csv?dl=1", SubscriptionTransportMedium.RemoteFile)

    def Reader(self, config, line, date, isLiveMode):
        if not (line.strip() and line[1].isdigit()): return None
        index = BRENT()
        index.Symbol = config.Symbol
        try:
            # Example File Format: (Data starts from 08/11/2008)
            # date     value
            # 8/11/08    110.54
            data = line.split(',')
            index.Time = datetime.strptime(data[0], "%Y-%m-%d")
            index.Value = Decimal(data[1])
        except:
            return None
            
        return index