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)

##### Download data#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','DAL']

#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 = 'DAL'
dfs[ticker].tail(5)

def bull_engulfing(c: pd.Series, o: pd.Series)-> pd.Series:
    """
    c: pd.Series. close price
    o: pd.Series. open price
    
    BULL ENGULFING: meet all of the following
        current OPEN > current CLOSE
        pre CLOSE > pre OPEN
        current CLOSE > pre OPEN
        current OPEN > pre CLOSE
    """
    be = (c>o) & (c.shift(1)<o.shift(1)) & (c>o.shift(1)) & (o<c.shift(1))
    return pd.Series(be, name='BULL_ENGULF')


def bbands_trigger(c:pd.Series, l:pd.Series, bb_lower:pd.Series) -> pd.Series:
    """
    c: pd.Series. close price
    l: pd.Series. low price
    bb_lower: lower band of Bollinger bands
    
    Bollinger Bands Trigger: meet all of the following
        pre CLOSE < pre BBANDS lower band
        current CLOSE > current BBANDS lower band
        current LOW < current BBANDS lower band
    """
    bt = (c.shift(1)<bb_lower.shift(1)) & (c > bb_lower) & (l < bb_lower)
    
    return pd.Series(bt, name='BBANDS_TRIGGER')

from core.finta import TA

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

TA.BBANDS

df_ta = TA.BBANDS(df, period = 20, std_multiplier=2.2, column="close")
df = df.merge(df_ta, left_index = True, right_index = True, how='inner' )


del df_ta
gc.collect()

df['BULL_ENGULF'] = bull_engulfing(df["Close"], df["Open"])
df['BBANDS_TRIGGER'] = bbands_trigger(df["Close"], df["Low"], df["BB_LOWER"])

df['SIGNAL'] = (df['BULL_ENGULF'] & df['BBANDS_TRIGGER']).astype(int)

df['B'] = df['SIGNAL']*(df["High"] + df["Low"])/2

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

df['SIGNAL'].value_counts()

from core.visuals import *

start = -1250
end = df.shape[0]
df_sub = df.iloc[start:end]
# df_sub = df[(df.index<='2019-04-01') & (df.index>='2019-01-24')]
names = {'main_title': f'{ticker}'}

lines0 = basic_lines(df_sub[['BB_UPPER', 'BB_MIDDLE', 'BB_LOWER']], 
                     colors = [], 
                     **dict(panel=0, width=1.5, secondary_y=False))

lines1 = basic_lines(df_sub[['BBWIDTH']], 
                     colors = ['cadetblue'], 
                     **dict(panel=1, width=1))
lines3 = basic_lines(df_sub[['PERCENT_B']], 
                     colors = ['lightcoral'], 
                     **dict(panel=2, width=1))

lines2 = basic_lines(df_sub[[ 'B']],
                     colors = ['navy'], 
                     **dict(panel=0, type='scatter', marker=r'${B}$' , markersize=100, secondary_y=False))


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

#shadows_ = basic_shadows(bands=[-0.01, 0.01], nsamples=df.iloc[start:end].shape[0], **dict(panel=1, color="lightskyblue",alpha=0.1,interpolate=True))
shadows_ = []
fig_config_ = dict(figratio=(18,10), volume=False, volume_panel=2,panel_ratios=(4,2, 2), tight_layout=True, returnfig=True,)

ax_cfg_ = {0:dict(basic=[5, 2, ['BB_UPPER', 'BB_MIDDLE', 'BB_LOWER']], 
                 title=dict(label = 'MA', fontsize=9, style='italic',  loc='left'), 
                ),
           2:dict(basic=[1, 0, ['BBWIDTH']]
                ),
           4:dict(basic=[1, 0, ['PERCENT_B']]
                ),
          }


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

aa_, bb_ = make_panels(main_data = df_sub[['Open', 'High', 'Low', 'Close', 'Volume']], 
                       added_plots = lines_,
                       fill_betweens = shadows_, 
                       fig_config = fig_config_, 
                       axes_config = ax_cfg_,  
                       names = names)

TRADE_CONFIG = dict(INIT_CAPITAL = 10000 ,
                    MIN_TRADE_SIZE = 100 ,
                    MAX_TRADE_SIZE = 1000 ,
                    HOLD_DAYS = 40, #max hold days
                    STOP_LOSS = 0.085, #10% drop
                    KEEP_PROFIT = 0.065, 
                    MAX_OPEN = 1, #allow only 1 open position
                    COST = 0.0035,
                   ) 

df['SIGNAL'].value_counts()

trades = []
for i in range(df.shape[0]-5):
    row = df.iloc[i]
    if row['SIGNAL']>0:
        print('enter: ', i)
        row_j = df.iloc[i+1]
        item = dict(signal_date = row.name,
                    enter_date = row_j.name, 
                    enter_price = row_j['High']
                   )
        for j in range(i+2, min(i+TRADE_CONFIG['HOLD_DAYS'], df.shape[0])):
            row_j = df.iloc[j]
            price_ = row_j['Low']
            pct_chg = price_/item['enter_price']
            if (pct_chg<= (1 - TRADE_CONFIG['STOP_LOSS'])) | (pct_chg >= (1 + TRADE_CONFIG['KEEP_PROFIT'])):
                break
                
        item['exit_date'] = row_j.name
        item['exit_price'] = price_
        item['hold_days'] = j - i
        i = j 
        print('exit:', i)
        trades.append(item)
        

df_trades = pd.DataFrame(data = trades)
df_trades.shape

def cal_pnl(trade):
    shares = int(TRADE_CONFIG['INIT_CAPITAL']/trade['enter_price'])
    if shares < TRADE_CONFIG['MIN_TRADE_SIZE']:
        shares = 0
    elif shares > TRADE_CONFIG['MAX_TRADE_SIZE']:
        shares = TRADE_CONFIG['MAX_TRADE_SIZE']
    pnl = shares*(trade['exit_price'] - trade['enter_price']) - shares*trade['enter_price']*TRADE_CONFIG['COST']
    return pnl

df_trades['pnl'] = df_trades.apply(lambda x: cal_pnl(x), axis=1)

df_trades