import warnings
import quandl
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import norm, t
from statsmodels.tsa.stattools import adfuller


pd.set_option('display.max_rows', 30)
plt.style.use('bmh')
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 15
warnings.filterwarnings('ignore')
quandl.ApiConfig.api_key = QUANDL_API_KEY

# exploratory settings
symbols = ('DRIP', 'XOP')
dates = ('2015-12-02', '2016-12-01')
M = 5

# retrieve data
df_X = quandl.get('EOD/' + symbols[0], start_date=dates[0], end_date=dates[1])
df_Y = quandl.get('EOD/' + symbols[1], start_date=dates[0], end_date=dates[1])
X = df_X.Adj_Close
Y = df_Y.Adj_Close

# calculate M-day returns
rX = (X.diff(M) / X.shift(M)).dropna()
rY = (Y.diff(M) / Y.shift(M)).dropna()
rY3 = rY * -3
spread = rX - rY3

# concat all variables
df = pd.concat([X, Y, rX, rY, rY3, spread], axis=1).dropna()
del X, Y, rX, rY, rY3, spread, df_X, df_Y
df.columns = [symbols[0], symbols[1],
              'r' + symbols[0], 'r' + symbols[1],
              'r' + symbols[1] + 'n3', 'spread']

df.head(10)

fig = plt.figure(figsize=(20, 7.5))

ax = fig.add_subplot(121)
ax.hist(df.spread, bins=50, edgecolor='k', facecolor='c', density=True)
x_grid = np.linspace(-.2, .2, 200)

# fit to normal
params = norm.fit(df.spread)
pdf_fitted = norm.pdf(x_grid, *params)
plt.plot(x_grid, pdf_fitted, color='r',
         label='$\mathcal{N}$' + '({:.2f},{:.2f})'.format(*params))

# fit to t
params = t.fit(df.spread)
pdf_fitted = t.pdf(x_grid, *params)
plt.plot(x_grid, pdf_fitted, color='orange',
         label='t({:.2f},{:.2f},{:.2f})'.format(*params))

ax.set_xlabel('Spread')
ax.set_ylabel('Density')
ax.legend(frameon=False, loc='upper left')

ax = fig.add_subplot(122)
df.spread.plot(ax=ax, color='k')
plt.xticks(rotation=0, ha='center')
ax.set_xlabel('Date')
ax.set_ylabel('Spread')

plt.tight_layout()
plt.show()

result = adfuller(df.spread)
print('ADF Statistic: {}\np-value: {}\nCritical Values:'.format(*result[:2]))
for key, value in result[4].items():
    print('\t{}: {:.3f}'.format(key, value))

def hurst(ts):
    lags = range(2, 100)
    tau = [np.sqrt(np.std(np.subtract(ts[lag:], ts[:-lag]))) for lag in lags]
    poly = np.polyfit(np.log(lags), np.log(tau), 1)
    return poly[0]*2.0

np.random.seed(123)
gbm = np.log(np.cumsum(np.random.randn(100000))+1000)
mr = np.log(np.random.randn(100000)+1000)
tr = np.log(np.cumsum(np.random.randn(100000)+1)+1000)

print('H: {:.4f}'.format(hurst(df.spread)))

be = BacktestEngine('DRIP', 'XOP', '2016-12-02', '2018-12-31', ratio=-3)
be.run(M=5, g=0.010, j=0.005, s=0.010, verbose=True)  # plot the backtest results

class BacktestEngine:
    
    def __init__(self, X_symbol, Y_symbol, start_date, end_date, ratio=1):
        '''
        PARAMETER:
            X_symbol: string
            Y_symbol: string
            start_date: string or datetime object
            end_date: string or datetime object

        RETURN: None
        '''

        self.symbols = (X_symbol, Y_symbol)
        self.__dates = (start_date, end_date)
        self.__stop_loss_dates = []
        self.__ratio = ratio
        
        # load data
        df_X = quandl.get('EOD/' + self.symbols[0],
                          start_date=self.__dates[0],
                          end_date=self.__dates[1])
        df_Y = quandl.get('EOD/' + self.symbols[1],
                          start_date=self.__dates[0],
                          end_date=self.__dates[1])
        self.__data = [df_X, df_Y]
        self.__trading = False

    def __process_data(self):  # prepare data to feed model
        '''
        PARAMETER: None
        
        RETURN: None
        '''
        
        df_X, df_Y = self.__data
        
        # generate price and volume series
        N_X = df_X.Adj_Volume * df_X.Adj_Close
        N_Y = df_Y.Adj_Volume * df_Y.Adj_Close
        N = pd.concat([N_X, N_Y], axis=1).min(axis=1).rolling(15).median()
        self.__K = N.max() * 2
        X = df_X.Adj_Close
        Y = df_Y.Adj_Close
        
        # calculate M-day returns
        M = self.__M
        rX = (X.diff(M) / X.shift(M)).dropna()
        rY = (Y.diff(M) / Y.shift(M)).dropna() * self.__ratio
        spread = rX - rY
        
        # concat all variables
        df = pd.concat([X, Y, rX, rY, spread, N], axis=1).dropna()
        del X, Y, rX, rY, spread, N
        df.columns = ['X', 'Y', 'rX', 'rY', 'spread', 'N']
        self.df = df
    
    def __update_params(self, M, g, j, s):  # update parameters of the model
        '''
        PARAMETER:
            M: integer
            g: float
            j: float
            s: float
        
        RETURN: None
        '''

        assert isinstance(M, (int, np.integer)) and (M > 0), \
               'parameter M should be a positive integer'
        assert -g < j < g, 'parameter j and g should follow -g < j < g'
        assert 0 < s, 'parameter s should be positive'

        self.__M = M
        self.__g = g
        self.__j = j
        self.__s = s
        self.__process_data()
    
    @staticmethod
    def __posval(position, price1, price2, gross=False):  # calculate value
        '''
        PARAMETER:
            position: float
            price1: float
            price2: float
            grodd: boolean (default: False)
        
        RETURN: float
        '''

        if gross:
            return abs(position[0]) * price1 + abs(position[1]) * price2
        else:
            return position[0] * price1 + position[1] * price2

    def __pos(self, value, X, Y, longX=True):  # calculate positions
        '''
        PARAMETER:
            value: float
            X: float
            Y: float
            longX: boolean (default: True)
        
        RETURN: (float, float)
        '''

        pX = round(value / X) * abs(self.__ratio)
        pY = round(value / Y) * (1 if self.__ratio > 0 else -1)
        return (pX, -pY) if longX else (-pX, pY)
    
    @staticmethod
    def __sharpe_ratio(daily_rtn):  # calculate sharpe ratio
        '''
        PARAMETER:
            daily_rtn: pd.Series of float
        
        RETURN: float
        '''
        
        return daily_rtn.mean() / daily_rtn.std()
    
    @staticmethod
    def __sortino_ratio(daily_rtn):  # calculate sortino ratio
        '''
        PARAMETER:
            daily_rtn: pd.Series of float
            
        RETURN: float
        '''
        daily_rtn = daily_rtn.copy()
        mean_rtn = daily_rtn.mean()
        daily_rtn.loc[daily_rtn > 0] = 0
        downside_risk = daily_rtn.std()
        return mean_rtn / downside_risk 
    
    @staticmethod
    def __maximum_drawdown(value):  # maximum drawdown
        '''
        PARAMETER:
            value: pd.Series of float
        
        RETURN: (float, datetime, datetime)
        '''
        
        end = (value.cummax() / value).idxmax()
        start = (value.loc[:end]).idxmax()
        mdd = 1 - value[end] / value[start]
        return mdd, start, end

    def run(self, M, g, j, s, verbose=True):  # run backtest using parameters
        '''
        PARAMETER:
            M: integer
            g: float
            j: float
            s: float
            verbose: boolean (default: True)
        
        RETURN: None
        '''
        
        self.__update_params(M, g, j, s)
        
        pX = []
        pY = []  # positions
        rtn = []
        curr_position = (0, 0)
        curr_rtn = 0
        prev_day = 0
        prev_cost = 0
        prev_gross_cost = 0
        spread = 0
        for date in self.df.index:
            if date.day < prev_day:  # check start of new month
                daily_rtn = curr_rtn
                curr_position = (0, 0)
                prev_cost = 0
                prev_gross_cost = 0
                self.__trading = True

            elif self.__trading:  # trading (never stop-lossed this month)
                if curr_rtn < -self.__s * prev_gross_cost:  # stop loss
                    self.__stop_loss_dates.append(date)
                    daily_rtn = curr_rtn
                    curr_position = (0, 0)
                    prev_cost = prev_gross_cost = 0
                    self.__trading = False

                else:
                    if curr_position[0] > 1:  # currently long in X
                        daily_rtn = curr_rtn
                        if spread > self.__g:  # direactly from long to short
                            curr_position = self.__pos(N / 100, X, Y, False)
                            prev_cost = self.__posval(curr_position, X, Y)
                            prev_gross_cost = self.__posval(curr_position, X, Y, True)
                        elif spread > -self.__j:  # close position
                            curr_position = (0, 0)
                            prev_cost = prev_gross_cost = 0
                    elif curr_position[0] < -1:  # currently short in X
                        daily_rtn = curr_rtn
                        if spread < -self.__g:  # directly from short to long
                            curr_position = self.__pos(N / 100, X, Y, True)
                            prev_cost = self.__posval(curr_position, X, Y)
                            prev_gross_cost = self.__posval(curr_position, X, Y, True)
                        elif spread < self.__j:  # close position
                            curr_position = (0, 0)
                            prev_cost = prev_gross_cost = 0
                    else:  # currently having no position
                        daily_rtn = 0
                        if spread > self.__g:  # open short position
                            curr_position = self.__pos(N / 100, X, Y, False)
                            prev_cost = self.__posval(curr_position, X, Y)
                            prev_gross_cost = self.__posval(curr_position, X, Y, True)
                        elif spread < -self.__g:  # open long position
                            curr_position = self.__pos(N / 100, X, Y, True)
                            prev_cost = self.__posval(curr_position, X, Y)
                            prev_gross_cost = self.__posval(curr_position, X, Y, True)

            else:  # ceased until next month
                daily_rtn = 0
            
            spread, X, Y, N = self.df[['spread', 'X', 'Y', 'N']].loc[date]
            curr_rtn = self.__posval(curr_position, X, Y) - prev_cost

            pX.append(curr_position[0])
            pY.append(curr_position[1])
            rtn.append(daily_rtn)
            prev_day = date.day

        self.df['pX'] = pX
        self.df['pY'] = pY
        self.df['daily_rtn'] = rtn
        self.df.daily_rtn /= self.__K  # rtn is stored in %
        self.df['cum_rtn'] = self.df.daily_rtn.cumsum()
        sortino = self.__sortino_ratio(self.df.daily_rtn)
        sharpe = self.__sharpe_ratio(self.df.daily_rtn)
        mdd = self.__maximum_drawdown(self.df.cum_rtn + self.__K)[0]
        yoy_rtn = (1 + self.df.cum_rtn[-1])**(250 / self.df.shape[0]) - 1
        if verbose: self.__plot()
        return sortino, sharpe, mdd, yoy_rtn
 
    def __plot(self):  # plot backtest results
        '''
        PARAMETER: None
        
        RETURN: None
        '''
        
        fig = plt.figure(figsize=(20, 15))

        ax1 = fig.add_subplot(221)
        self.df.rX.plot(ax=ax1, label='X', alpha=.5, color='darkgreen')
        self.df.rY.plot(ax=ax1, label='Y', color='pink')
        ax1.set_ylabel('Returns')
        self.df.spread.plot(ax=ax1, label='spread', color='k')
        plt.xticks(rotation=0, ha='center')
        plt.legend(frameon=False, loc='upper right')
        
        ax2 = fig.add_subplot(222)
        self.df.spread.plot(ax=ax2, label='spread', color='k')
        ax2.plot((min(self.df.index), max(self.df.index)),
                 (self.__g, self.__g), 'c--', label='$\pm$g')
        ax2.plot((min(self.df.index), max(self.df.index)),
                 (self.__j, self.__j), 'r--', label='$\pm$j')
        ax2.plot((min(self.df.index), max(self.df.index)),
                 (-self.__g, -self.__g), 'c--')
        ax2.plot((min(self.df.index), max(self.df.index)),
                 (-self.__j, -self.__j), 'r--')
        ax2.set_ylabel('Spread')
        plt.xticks(rotation=0, ha='center')
        plt.legend(frameon=False)
        
        mdd, sd, ed = self.__maximum_drawdown(self.df.cum_rtn + self.__K)
        
        ax3 = fig.add_subplot(223)
        self.df.N.plot(ax=ax3, label='N', color='k')
        ax3.set_ylabel('Total Dollar Volumes')
        plt.legend(frameon=False, loc='upper left')
        plt.xticks(rotation=0, ha='center')
        ax3_ = ax3.twinx()
        self.df.pX.abs().plot(ax=ax3_, label='position', alpha=.5, color='r')
        ax3_.ticklabel_format(axis='y', style='sci', scilimits=(5, 5))
        ax3_.axvspan(xmin=sd, xmax=ed, ymin=0, ymax=1, color='grey', alpha=.5)
        ax3_.set_ylim(ymin=0)
        ax3_.set_ylabel('Daily Position (Long)')
        ax3.grid(False)
        plt.legend(frameon=False, loc='upper right')
        
        ax4 = fig.add_subplot(224)
        ax4.axvspan(xmin=sd, xmax=ed, ymin=0, ymax=1, color='grey', alpha=.5)
        (self.df.daily_rtn * 100).plot(ax=ax4, label='Daily Return (%)', color='c')
        (self.df.cum_rtn * 100).plot(ax=ax4, label='Cum. Return (%)', color='r')
        ax4.scatter(self.__stop_loss_dates,
                    self.df.cum_rtn.loc[self.__stop_loss_dates] * 100,
                    marker='x', color='k', label='', zorder=10)
        ax4.set_ylabel('Return (%)')
        plt.xticks(rotation=0, ha='center')
        plt.legend(frameon=False)
        
        plt.tight_layout()
        plt.show()

be = BacktestEngine('DRIP', 'XOP', start_date='2016-12-02', end_date='2018-12-31', ratio=-3)
st, sr, md, rt = be.run(M=5, g=0.02, j=0.01, s=0.01, verbose=True)
print('Sortino Ratio={:.4f}, Sharpe Ratio={:.4f}, Maximum Drawdown={:.4g}, YoY Return={:.2%}'
      .format(st, sr, md ,rt))

be.df.loc[be.df.pX != 0]

from time import time
from itertools import product


def record():
    # record function with closure designed
    # to keep track of the best parameters
    record.best = []
    def f(*args):
        record.best.append((*args,))
    return f

mst, msr, mmd, mrt = 0, 0, 1, 0
M_grid = np.arange(5, 25, 5)
g_grid = np.arange(0.001, 0.012, 0.002)
j_grid = np.arange(-0.010, 0.011, 0.002)
s_grid = [1e-3, 5e-3, 0.01, 0.05, 0.10]
nsim = len(M_grid) * len(g_grid) * len(j_grid) * len(s_grid)
report_params = record()

be = BacktestEngine('DRIP', 'XOP', start_date='2015-12-02', end_date='2016-12-01', ratio=-3)

start_time = time()
for i, params in enumerate(product(M_grid, g_grid, j_grid, s_grid)):
    try:
        st, sr, md, rt = be.run(*params, verbose=False)
        if np.isnan(sr) or sr < 0 or st < 0: continue
        report = False
        if st > mst: mst, report = st, True
        if sr > msr: msr, report = sr, True
        if md < mmd: mmd, report = md, True
        if rt > mrt: mrt, report = rt, True
        prog = (i + 1) / nsim
        current_time = time()
        elapse = current_time - start_time
        eta = elapse / prog * (1 - prog)
        print('Running {} simulations... {:.2%} finished, ETA={:.0f}s'
              .format(nsim, prog, eta) + ' ' * 10, end='\r')
        if report: report_params(*params, st, sr, md, rt)
    except AssertionError:  # (-g < j < g) does not hold
        continue

st_list, sr_list, md_list, rt_list = list(zip(*record.best))[-4:]
record.best = [entries for _, entries in sorted(zip(st_list, record.best))][-5:]
record.best = [entries for _, entries in sorted(zip(sr_list, record.best))][-5:]
record.best = [entries for _, entries in sorted(zip(rt_list, record.best))][-5:]
record.best = [entries for _, entries in sorted(zip(md_list, record.best))][:5]

for i, entries in enumerate(record.best):
    print('(Record {:2}) M={:>2}, g={:.3f}, j={:6.3f}, s={:.5f},\t'
          'st={:.4f} sr={:.4f}, md={:9.4g}, rt={:6.2%}'
          .format(i, *entries))

rec = np.arange(len(record.best))
st_list, st_list, md_list, rt_list = list(zip(*record.best))[-4:]

fig = plt.figure(figsize=(20, 7.5))

ax = fig.add_subplot(121)
ax.set_xlim(min(rt_list) * 0.95, max(rt_list) * 1.05)
ax.set_ylim(min(st_list) * 0.95, max(st_list) * 1.05)
st_offset = (max(st_list) - min(st_list)) * 0.02
ax.scatter(rt_list, st_list, s=[(md / np.mean(md_list))**2 * 4e3 for md in md_list],
           color='r', edgecolor='k', alpha=.5)
ax.text(.086, .325, '(Bubble Size ~ Maximum Drawdown)')
for rt, st, i in zip(rt_list, st_list, range(len(record.best))):
    ax.text(x=rt, y=st - st_offset, s=str(i), ha='center')
ax.set_xlabel('YoY Return')
ax.set_ylabel('Sortino Ratio')

ax = fig.add_subplot(122)
ax.scatter(rec, st_list, color='r', marker='x', s=200, label='Sortino Ratio')
ax.scatter(rec, rt_list, color='c', marker='o', s=200, label='YoY Return')
ax.set_xlabel('Record')
ax.set_ylabel('Return & Ratio')
ax.legend(frameon=False, loc='upper left')
 
ax = ax.twinx()
ax.grid(False)
ax.scatter(rec, md_list, color='orange', marker='s', s=200, label='Maximum Drawdown')
ax.set_ylim(min(md_list) * 0.9, max(md_list) * 1.1)
ax.xaxis.set_ticks(rec)
ax.set_ylabel('Maximum Drawdown')
ax.legend(frameon=False, loc='upper right')

plt.tight_layout()

be = BacktestEngine('DRIP', 'XOP', start_date='2016-12-02', end_date='2018-12-31', ratio=-3)
st, sr, md, rt = be.run(M=20, g=0.007, j= 0.006, s=0.05000, verbose=True)
print('Sortino Ratio={:.4f}, Sharpe Ratio={:.4f}, Maximum Drawdown={:.4g}, YoY Return={:.2%}'
      .format(st, sr, md ,rt))

be = BacktestEngine('DRIP', 'XOP', start_date='2016-12-02', end_date='2018-12-31', ratio=-3)
st, sr, md, rt = be.run(M=20, g=0.007, j=-0.006, s=0.05000, verbose=True)
print('Sortino Ratio={:.4f}, Sharpe Ratio={:.4f}, Maximum Drawdown={:.4g}, YoY Return={:.2%}'
      .format(st, sr, md ,rt))