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)

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

benchmark_tickers = ['^GSPC']
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'
dfs[ticker].tail(5)

import math
def cal_rsih(ohlc: pd.DataFrame, 
               period: int = 14,
               column: str = "close") -> pd.Series:
    """
    source: https://traders.com/Documentation/FEEDbk_docs/2022/01/TradersTips.html
    // TASC JAN 2022, RSIH - RSI with Hann Windowing, John F. Ehlers

    inputs:
        RSILength(14);

    // Accumulate "Closes Up" and "Closes Down"
    CU = 0;
    CD = 0;

    for count = 1 to RSILength begin
        if Close[count - 1] - Close[count] > 0 then 
         CU = CU + (1 - Cosine(360*count / (RSILength + 1)))
         *(Close[count - 1] - Close[count]);
        if Close[count] - Close[count - 1] > 0 then 
         CD = CD + (1 - Cosine(360*count / (RSILength + 1)))
         *(Close[count] - Close[count - 1]);
    end;

    if CU + CD <> 0 then 
        MyRSI = (CU - CD) / (CU + CD);

    """
    
    def _hann(c, rsi_len):
        cu = 0
        cd = 0
    
        for i in range(1, rsi_len):
            j = i + 1
            delta = c[i] - c[i-1] #e.g. i=1, delta = c[0] - c[-1] => current close minus previous close. 
            if delta>0:
                cu = cu + (1 - math.cos(360*j/(rsi_len + 1)))*delta
            else:
                cd = cd - (1 - math.cos(360*j/(rsi_len + 1)))*delta
        re = 0
        if (cu + cd) != 0:
            re = (cu - cd)/(cu + cd)
        return re
    
    
    c = ohlc[column]
    rsi_ = c.rolling(window=period, min_periods=period).apply(lambda x: _hann(x, period))
    
    return pd.Series(rsi_,index=ohlc.index, name=f"RSIH")    

import math
def cal_rsih(ohlc: pd.DataFrame, 
               period: int = 14,
               column: str = "close") -> pd.Series:
    """
    source: https://traders.com/Documentation/FEEDbk_docs/2022/01/TradersTips.html
    
    lengthInput = input.int(14, "Length:", minval = 2)

    rsih(length) =>
        var float PIx2 = 2 * math.pi
        // Accumulate "Closes Up" and "Closes Down"
        cu = 0.0
        cd = 0.0
        for count = 1 to length
            change = close[count] - close[count - 1]
            absChange = math.abs(change)
            cosPart = math.cos(PIx2 * count / (length + 1))

            if change < 0
                cu := cu + (1 - cosPart) * absChange
            else if change > 0
                cd := cd + (1 - cosPart) * absChange

        result = nz((cu - cd) / (cu + cd))
    """
 
    def _hann(_data, _len):
        pi_ = 2*math.pi
        
        #Accumulate "Closes Up" and "Closes Down"
        cu = 0.0
        cd = 0.0
    
        for i in range(1, _len):
            delta = _data[i] - _data[i-1] 
            delta_abs = np.abs(delta)
            cos_ = math.cos(pi_*i/(_len + 1))
            
            if delta>0:
                cu = cu + (1 - cos_)*delta_abs
            else:
                cd = cd + (1 - cos_)*delta_abs
        re = 0
        if (cu + cd) != 0:
            re = (cu - cd)/(cu + cd)
        return re
    
    
    c = ohlc[column]
    rsi_ = c.rolling(window=period, min_periods=period).apply(lambda x: _hann(x, period))
    

    return pd.Series(rsi_,index=ohlc.index, name=f"RSIH")    
    
    

def _hann(_data, _len):
    out_ = np.zeros(_len)
    for i in range(_len):
        out_[i] = _data[i]*(1-math.cos(2*math.pi*(i+1)/(_len + 1)))

df = dfs[ticker][['Open', 'High', 'Low', 'Close', 'Volume']]

df = df.round(2)

cal_rsih

df_ta = cal_rsih(df, period=14, column="Close")
df = df.merge(df_ta, left_index = True, right_index = True, how='inner' )

del df_ta
gc.collect()

from core.finta import TA

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

del df_ta
gc.collect()

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

df[['RSIH']].hist(bins=50)

#https://github.com/matplotlib/mplfinance
#this package help visualize financial data
import mplfinance as mpf
import matplotlib.colors as mcolors

# all_colors = list(mcolors.CSS4_COLORS.keys())#"CSS Colors"
# all_colors = list(mcolors.TABLEAU_COLORS.keys()) # "Tableau Palette",
# all_colors = list(mcolors.BASE_COLORS.keys()) #"Base Colors",
all_colors = ['dodgerblue', 'firebrick','limegreen','skyblue','lightgreen',  'navy','yellow','plum',  'yellowgreen']
#https://github.com/matplotlib/mplfinance/issues/181#issuecomment-667252575
#list of colors: https://matplotlib.org/stable/gallery/color/named_colors.html
#https://github.com/matplotlib/mplfinance/blob/master/examples/styles.ipynb

def make_3panels2(main_data, add_data, mid_panel=None, chart_type='candle', names=None, figratio=(14,9)):

    style = mpf.make_mpf_style(base_mpf_style='yahoo',  #charles
                               base_mpl_style = 'seaborn-whitegrid',
#                                marketcolors=mpf.make_marketcolors(up="r", down="#0000CC",inherit=True),
                               gridcolor="whitesmoke", 
                               gridstyle="--", #or None, or - for solid
                               gridaxis="both", 
                               edgecolor = 'whitesmoke',
                               facecolor = 'white', #background color within the graph edge
                               figcolor = 'white', #background color outside of the graph edge
                               y_on_right = False,
                               rc =  {'legend.fontsize': 'small',#or number
                                      #'figure.figsize': (14, 9),
                                     'axes.labelsize': 'small',
                                     'axes.titlesize':'small',
                                     'xtick.labelsize':'small',#'x-small', 'small','medium','large'
                                     'ytick.labelsize':'small'
                                     }, 
                              )   

    if (chart_type is None) or (chart_type not in ['ohlc', 'line', 'candle', 'hollow_and_filled']):
        chart_type = 'candle'
    len_dict = {'candle':2, 'ohlc':3, 'line':1, 'hollow_and_filled':2}    
        
    kwargs = dict(type=chart_type, figratio=figratio, volume=True, volume_panel=2, 
                  panel_ratios=(4,2, 2), tight_layout=True, style=style, returnfig=True)
    
    if names is None:
        names = {'main_title': '', 'sub_tile': ''}
    
    added_plots = { 
#         'S':  mpf.make_addplot(add_data['S'], panel=0, color='blue', type='scatter', marker=r'${S}$' , markersize=100, secondary_y=False),   
#         'B':  mpf.make_addplot(add_data['B'], panel=0, color='blue', type='scatter', marker=r'${B}$' , markersize=100, secondary_y=False), 
        
        'EMA': mpf.make_addplot(add_data['EMA'], panel=0, color='dodgerblue', width=1 ,secondary_y=False), 
    }

        

    if mid_panel is not None:
        i = 0
        for name_, data_ in mid_panel.iteritems():
            added_plots[name_] = mpf.make_addplot(data_, panel=1, width=1,color=all_colors[i])
            i = i + 1
#         fb_bbands2_ = dict(y1=-0.5*np.ones(mid_panel.shape[0]),
#                       y2=0.5*np.ones(mid_panel.shape[0]),color="lightskyblue",alpha=0.1,interpolate=True)
#         fb_bbands2_['panel'] = 1
#         fb_bbands.append(fb_bbands2_)
     
        
    fb_bbands = []
    fb_span_up = dict(y1=np.zeros(mid_panel.shape[0]),y2=mid_panel['RSIH'].values,where=mid_panel['RSIH']<0,color="#FF008055",alpha=0.2, panel=1, interpolate=True)
    fb_span_dn = dict(y1=np.zeros(mid_panel.shape[0]),y2=mid_panel['RSIH'].values,where=mid_panel['RSIH']>0,color="palegreen",alpha=0.2, panel=1, interpolate=True)


    fb_bbands= [fb_span_up, fb_span_dn]
    

    fig, axes = mpf.plot(main_data,  **kwargs,
                         addplot=list(added_plots.values()), 
                         fill_between = fb_bbands,
                        )
    # add a new suptitle
    fig.suptitle(names['main_title'], y=1.05, fontsize=12, x=0.1285)

#     axes[0].legend([None]*4)
#     handles = axes[0].get_legend().legendHandles
#     axes[0].legend(handles=handles[2:],labels=['RS_EMA', 'EMA'])
#     axes[0].set_title(names['sub_tile'], fontsize=10, style='italic',  loc='left')
    

#     axes[0].set_ylabel(names['y_tiles'][0])
#     axes[2].set_ylabel(names['y_tiles'][1])
    return fig, axes
   

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

names = {'main_title': f'{ticker}', 
         'sub_tile': 'RSI with Hann'}


aa_, bb_ = make_3panels2(df.iloc[start:end][['Open', 'High', 'Low', 'Close', 'Volume']], 
            df.iloc[start:end][['EMA']],
            df.iloc[start:end][['RSIH']],
             chart_type='hollow_and_filled',names = names)