
import pandas_datareader as pdr
import datetime


ticker = 'TSLA'
ohlcv = pdr.get_data_yahoo(ticker, datetime.date.today() - datetime.timedelta(1025), datetime.date.today())

def ATR(DF, n):
    '''
    Function to calculate True Range and Average True Range
    
    Parameters
    ----------
    DF : DataFrame
        DESCRIPTION.
    n : int
        number of days.

    Returns
    -------
    DataFrame.

    '''
    df = DF.copy()
    df['H-L'] =  abs(df['High'] - df['Low'])
    df['H-PC'] = abs(df['High'] - df['Adj Close'].shift(1))
    df['L-PC'] = abs(df['Low'] - df['Adj Close'].shift(1)) # shift(n): Move n positive direction
    df['TR'] = df[['H-L', 'H-PC', 'L-PC']].max(axis=1, skipna=False)
    df['ATR'] = df['TR'].rolling(n).mean() # rolling(n): Take window of n
    # df['ATR'] = df['TR'].ewm(span=n, adjust=False, min_periods=n).mean() # Exponential mean
    
    df2 = df.drop(['H-L', 'H-PC', 'L-PC'], axis=1)
    return df2

res = ATR(ohlcv, 20)

def BollBand(DF, n):
    '''
    Function to caculate Bollinger Band

    Parameters
    ----------
    DF : DataFrame
        DESCRIPTION.
    n : int
        number of days.

    Returns
    -------
    DataFrame.

    '''
    
    df = DF.copy()
    df['MA'] = df['Adj Close'].rolling(n).mean()
    df['BB_up'] = df['MA'] + 2 * df['MA'].rolling(n).std()
    df['BB_dn'] = df['MA'] - 2 * df['MA'].rolling(n).std()
    df['BB_width'] = df['BB_up'] - df['BB_dn']
    df.dropna(inplace=True)
    return df

res2 = BollBand(ohlcv, 20).iloc[-100:, [-4,-3,-2]].plot()

    
    
    
    