import pandas as pd
import numpy as np
import os
import gc
import copy
from pathlib import Path
from datetime import datetime, timedelta, time, date

#this package is to download equity price data from yahoo finance
#the source code of this package can be found here: https://github.com/ranaroussi/yfinance/blob/main
import yfinance as yf

pd.options.display.max_rows = 100
pd.options.display.max_columns = 100

import warnings
warnings.filterwarnings("ignore")

import pytorch_lightning as pl
random_seed=1234
pl.seed_everything(random_seed)

#S&P 500 (^GSPC),  Dow Jones Industrial Average (^DJI), NASDAQ Composite (^IXIC)
#Russell 2000 (^RUT), Crude Oil Nov 21 (CL=F), Gold Dec 21 (GC=F)
#Treasury Yield 10 Years (^TNX)
#CBOE Volatility Index (^VIX) Chicago Options - Chicago Options Delayed Price. Currency in USD

#benchmark_tickers = ['^GSPC', '^DJI', '^IXIC', '^RUT',  'CL=F', 'GC=F', '^TNX']

benchmark_tickers = ['^GSPC', '^VIX']
tickers = benchmark_tickers + ['GSK', 'BST', 'PFE']

#https://github.com/ranaroussi/yfinance/blob/main/yfinance/base.py
#     def history(self, period="1mo", interval="1d",
#                 start=None, end=None, prepost=False, actions=True,
#                 auto_adjust=True, back_adjust=False,
#                 proxy=None, rounding=False, tz=None, timeout=None, **kwargs):

dfs = {}

for ticker in tickers:
    cur_data = yf.Ticker(ticker)
    hist = cur_data.history(period="max", start='2000-01-01')
    print(datetime.now(), ticker, hist.shape, hist.index.min(), hist.index.max())
    dfs[ticker] = hist

ticker = '^GSPC'
# ticker = 'PFE'
dfs[ticker].tail(5)

def _zscore(_data, _len):
    roll_ = _data.rolling(window=_len, min_periods = _len)
    z_ = (_data - roll_.mean())/roll_.std()
    return z_

from core.finta import TA

z_period = 10
fast_period = 10
slow_period = 100

help(TA.SMA)

df = dfs[ticker][['Open', 'High', 'Low', 'Close', 'Volume']]
df = df.round(2)

df_ta = TA.SMA(df, period = fast_period, column="close")
df_ta.name='MA_FAST'
df = df.merge(df_ta, left_index = True, right_index = True, how='inner' )

df_ta = TA.SMA(df, period = slow_period, column="close")
df_ta.name='MA_SLOW'
df = df.merge(df_ta, left_index = True, right_index = True, how='inner' )

del df_ta
gc.collect()

df['z_score'] = _zscore(df['Close'], z_period)

display(df.head(5))
display(df.tail(5))

df.dropna(how='any', inplace=True)

#1: BUY signal: z-score < -1 AND ma_fast > ma_slow
#2: ma_fast crosses over ma_slow: BUY to COVER
#-1: SELL SHORT signal: z-score > 1 AND ma_fast < ma_slow
#-2: ma_fast crosses under ma_slow: SELL
enter_zscore = 1
exit_zscore = 0.5

signals = [0]
for i in range(1, df.shape[0]):
    row = df.iloc[i]
    z_ = row['z_score']
    ma_f = row['MA_FAST']
    ma_s = row['MA_SLOW']
    
    row_j = df.iloc[i-1]
    pre_ma_f = row_j['MA_FAST']
    pre_ma_s = row_j['MA_SLOW']
     
    s_ = 0
    if (z_ < -enter_zscore) & (ma_f > ma_s):
        s_ = 1
    elif (z_ > enter_zscore) & (ma_f < ma_s):
        s_ = -1
    elif (pre_ma_f < pre_ma_s) & (ma_f >= ma_s):
        s_ = 2
    elif (pre_ma_f > pre_ma_s) & (ma_f <= ma_s):
        s_ = -2
    
    signals.append(s_)
    
df['SIGNAL'] = signals

df['SIGNAL'].value_counts()

df['B1'] = (df['SIGNAL']==1).astype(int)*(df['High'] + df['Low'])/2
df['B2'] = (df['SIGNAL']==2).astype(int)*(df['High'] + df['Low'])/2
df['S1'] = (df['SIGNAL']==-1).astype(int)*(df['High'] + df['Low'])/2
df['S2'] = (df['SIGNAL']==-2).astype(int)*(df['High'] + df['Low'])/2

from core.visuals import *

start = -250
end = df.shape[0]

names = {'main_title': f'{ticker}'}

lines0 = basic_lines(df.iloc[start:end][['MA_FAST', 'MA_SLOW']], 
                     colors = [], 
                     **dict(panel=0, width=1.5, secondary_y=False))

lines1 = basic_lines(df.iloc[start:end][['z_score']], 
                     colors = ['cadetblue'], 
                     **dict(panel=1, width=1, secondary_y=False))

lines2 = basic_lines(df.iloc[start:end][[ 'B2']],
                     colors = ['navy'], 
                     **dict(panel=0, type='scatter', marker=r'${B}$' , markersize=100, secondary_y=False))


lines_ = dict(**lines0, **lines1)
lines_.update(lines2)

shadows_ = basic_shadows(bands=[-1, 1], nsamples=df.iloc[start:end].shape[0], **dict(panel=1, color="lightskyblue",alpha=0.1,interpolate=True))

fig_config_ = dict(figratio=(18,10), volume=False, volume_panel=2,panel_ratios=(4,3), tight_layout=True, returnfig=True,)

ax_cfg_ = {0:dict(basic=[4, 2, ['MA_FAST', 'MA_SLOW']], 
                 title=dict(label = 'MA', fontsize=9, style='italic',  loc='left'), 
                ),
           2:dict(basic=[1, 0, ['z_score',]]
                ),
          }


names = {'main_title': f'{ticker}'}

aa_, bb_ = make_panels(main_data = df.iloc[start:end][['Open', 'High', 'Low', 'Close', 'Volume']], 
                       added_plots = lines_,
                       fill_betweens = shadows_, 
                       fig_config = fig_config_, 
                       axes_config = ax_cfg_,  
                       names = names)

TRADE_CONFIG = dict(InitialCapital = 10000 ,
                    LeverageToUse =  2 ,
                    MinTradeSize = 100 ,
                    MaxTradeSize = 1000 
                   )