from zeta.csv.day import DayCSVLoader
OHLCV_COLS = ['Open', 'High', 'Low', 'Close', 'Volume']

loader = DayCSVLoader()
df = pd.concat([loader.read_csv('000660')[OHLCV_COLS],
                loader.read_csv('MU')[OHLCV_COLS]], axis=1)
df.columns = ['Open', 'High', 'Low', 'Close', 'Volume', 'Open_', 'High_', 'Low_', 'Close_', 'Volume_']
df.fillna(method='ffill', inplace=True)
df.dropna(inplace=True)
df = df[df.index > datetime(2015, 1, 1)].copy()

# Price
price1 = df['Close']  # .pct_change().cumsum().iloc[1:]
price2 = df['Close_'] # .pct_change().cumsum().iloc[1:]
price1.name = '000660'
price2.name = 'MU'

# Z-Score Normalization (주가가 정규분포를 따른다고 가정함.. 사실 틀림 -> 잘못 쓰면 손실 크게 남)
norm_price1 = (price1 - price1.mean()) / price1.std()
norm_price2 = (price2 - price2.mean()) / price2.std()

# Percent Price
pct_price1 = price1.pct_change().cumsum().iloc[1:]
pct_price2 = price2.pct_change().cumsum().iloc[1:]

# Percent Price + Z-Score Normalization
pct_norm_price1 = price1.pct_change().cumsum().iloc[1:]
pct_norm_price2 = price2.pct_change().cumsum().iloc[1:]
pct_norm_price1 = (pct_norm_price1 - pct_norm_price1.mean()) / pct_norm_price1.std()
pct_norm_price2 = (pct_norm_price2 - pct_norm_price2.mean()) / pct_norm_price2.std()

from statsmodels.tsa.stattools import coint

stat, p_value, _ = coint(price1, price2)
corr = price1.corr(price2)

print(f'Cointegration Stat    : {stat:>8.4f}')
print(f'Cointegration P-Value : {p_value:>8.4f}')
print(f'Correlation           : {corr:>8.4f}')

x = pd.Series(np.random.normal(1, 1, 100).cumsum())
y = pd.Series(np.random.normal(2, 1, 100).cumsum())
k = pd.Series(np.random.normal(0, 1, 100)) + 20
l = k.copy()
l[0:10] = 30
l[10:20] = 10
l[20:30] = 30
l[30:40] = 10
l[40:50] = 30
l[50:60] = 10
l[60:70] = 30
l[70:80] = 10
l[80:90] = 30
l[90:100] = 10

corr1 = x.corr(y)
_, coint1, _ = coint(x, y)

corr2 = k.corr(l)
_, coint2, _ = coint(k, l)

# Visualization
fig, ax = plt.subplots(1, 2, figsize=(12, 4))
ax[0].plot(x)
ax[0].plot(y)
ax[0].set_title(f'Corr:{corr1:.2f} | Coint:{1-coint1:.2f}')
ax[1].plot(k)
ax[1].plot(l)
ax[1].set_title(f'Corr:{corr2:.2f} | Coint:{1-coint2:.2f}')

import statsmodels.api as sm
from statsmodels.tsa.stattools import adfuller

x = sm.add_constant(pct_norm_price1)  # exog == x
y = pct_norm_price2                   # endog == y
ols = sm.OLS(y, x).fit()
bias = ols.params['const']   # bias=y-intercept
beta = ols.params['000660']  # beta=slope

print(f'bias (y-intercept): {bias:>7.4f}')
print(f'beta (slope)      : {beta:>7.4f}')

# Calculate Residual <- residual(spread) 값이 stationary 해야지 x, y는 서로 cointegration하다고 할 수 있음
spread = pct_norm_price2 - beta * pct_norm_price1
spread.name = 'Spread'

# Stationarity Test. p-value값이 0.05 이하라면 residual은 stationary 가능성이 높다
p_value = adfuller(spread)[1]
print(f'AdFuller P-Value  : {p_value:.4f}')

# Visualization
fig, ax = plt.subplots(1, 2, figsize=(12, 4))
fig.set_tight_layout(True)
ax[0].plot(price1, label='SK Hynix')
ax[0].plot(price2, label='MU')
ax[0].plot(beta * price1, color='red', label='beta * SK Hynix')
ax[0].legend()

ax[1].plot(spread, label='residual')
ax[1].axhline(spread.mean(), color='red')
ax[1].legend()

# Eagle-Granger Method
x = sm.add_constant(pct_norm_price1)  # exog == x
y = pct_norm_price2                   # endog == y
ols = sm.OLS(y, x).fit()  
spread = pct_norm_price2 - beta * pct_norm_price1
lower, upper = spread.quantile([0.2, 0.8])

# SK Hynix 데이터가 좀 튀는 부분이 있어서 과거 데이터 삭제 함
x = price1.loc[price1.index > datetime(2018, 1, 1)]
y = price2.loc[price2.index > datetime(2018, 1, 1)]

# Ratio Method
ratio = y / x
lower, upper = ratio.quantile([0.2, 0.8])

stock_df = df[OHLCV_COLS + ['Close_']].dropna().rename(columns={'Close_': 'RefPrice'})
stock_df.tail()

def PCT(arr):
    arr = pd.Series(arr)
    return arr.pct_change().cumsum().fillna(method='bfill')

def ZScore(arr):
    arr = pd.Series(arr)
    return (arr - arr.mean()) / arr.std()

def Spread(x, y):
    x = ZScore(PCT(x))
    y = ZScore(PCT(y))
    
    x_ = sm.add_constant(x)  # exog == x
    ols = sm.OLS(y, x_).fit() 
    spread = y - beta * x
    spread.name = 'Spread'
    return spread


spread = Spread(stock_df['Close'], stock_df['RefPrice'])
lower, upper = spread.quantile([0.3, 0.7])

from backtesting import Backtest, Strategy

class PairStrategy(Strategy):
    mean = 0.2
    upper = -0.1
    lower = -0.05
    
    def init(self):
        price1 = self.data.Close      # 전날 종가
        price2 = self.data.RefPrice   # 금일 미국장 종가

        # 전일종가 - 미국종가
        self.spread = self.I(Spread, price1, price2, name='Spread')

    def next(self):
        price1 = self.data.Close[-1]
        price2 = self.data.RefPrice[-1]
        spread = self.spread
        
        if spread[-2] <= self.upper and spread[-1] >= self.upper:
            order = self.buy(size=0.8, limit=price1)
        elif self.spread[-2] >= self.mean and self.spread[-1] <= self.mean:
            self.position.close()

bt = Backtest(stock_df,
              PairStrategy,
              cash=100000,
              commission=0.0036,
              exclusive_orders=True)
result = bt.run()
fig = bt.plot(show_legend=False)
display(result)
display(result._trades)

import seaborn as sns

stats, heatmap, optimize_result = bt.optimize(mean=np.arange(-0.5, 0.5, 0.01),
                                              upper=np.arange(-0.5, 0.5, 0.01),
                                              lower=np.arange(-0.5, 0.5, 0.01),
                                              maximize='Sharpe Ratio',
                                              max_tries=200,
                                              random_state=0,
                                              method='skopt',
                                              return_heatmap=True,
                                              return_optimization=True)

display(heatmap.sort_values(ascending=False).head())
sns.heatmap(heatmap.groupby(['mean', 'upper']).mean().unstack()[::-1],
            cmap='viridis')