# https://quantpedia.com/strategies/trading-wti-brent-spread/
#
# A 20-day moving average of WTI/Brent spread is calculated each day. If the current spread value is above SMA 20 then we enter a short position
# in the spread on close (betting that the spread will decrease to the fair value represented by SMA 20). The trade is closed at the close of the
# trading day when the spread crosses below fair value. If the current spread value is below SMA 20 then we enter a long position betting that 
# the spread will increase and the trade is closed at the close of the trading day when the spread crosses above fair value.

from collections import deque
import numpy as np

class WTI_BRENT_Spread(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 6, 28)
        self.SetCash(100000)
        
        self.symbols = [
            "CME_CL1",  # Crude Oil Futures, Continuous Contract
            "ICE_B1"    # Brent Crude Oil Futures, Continuous Contract
        ]

        self.spread = deque(maxlen=2)
        
        # True -> Quantpedia data
        # False -> Quandl free data
        self.use_quantpedia_data = True
        
        if not self.use_quantpedia_data:
            self.symbols = ['CHRIS/' + x for x in self.symbols]
        
        for symbol in self.symbols:
            data = None
            if self.use_quantpedia_data:
                data = self.AddData(QuantpediaFutures, symbol, Resolution.Daily)
            else:
                data = self.AddData(QuandlFutures, symbol, Resolution.Daily)
            #data.SetLeverage(2)
        
    def OnData(self, data):
        symbol1 = self.symbols[0]
        symbol2 = self.symbols[1]
        
        if self.Securities.ContainsKey(symbol1) and self.Securities.ContainsKey(symbol2):
            price1 = self.Securities[symbol1].Price
            price2 = self.Securities[symbol2].Price
            if price1 != 0 and price2 != 0:
                spread = price1 - price2
                self.spread.append(spread)
        
        # MA calculation
        if len(self.spread) == self.spread.maxlen:
            spreads = [x for x in self.spread]
            spread_ma20 = np.average(spreads)
            
            current_spread = spreads[-1]
            
            if current_spread > spread_ma20:
                if self.Portfolio[symbol1].IsLong:
                    self.Liquidate(symbol1)
                
                if not self.Portfolio[symbol1].IsShort:
                    self.SetHoldings(symbol1, -1)
            elif current_spread < spread_ma20:
                if self.Portfolio[symbol1].IsShort:
                    self.Liquidate(symbol1)
                
                if not self.Portfolio[symbol1].IsLong:
                    self.SetHoldings(symbol1, 1)

# Quantpedia data
class QuantpediaFutures(PythonData):
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("http://data.quantpedia.com/backtesting_data/futures/{0}.csv".format(config.Symbol.Value), SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config, line, date, isLiveMode):
        data = QuantpediaFutures()
        data.Symbol = config.Symbol
        
        try:
            if not line[0].isdigit(): return None
            split = line.split(';')
            
            data.Time = datetime.strptime(split[0], "%d.%m.%Y") + timedelta(days=1)
            data['settle'] = float(split[1])
            data.Value = float(split[1])
        except:
            return None
            
        return data

# Quandl free data
class QuandlFutures(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = "settle"