# Mount Google colab
from google.colab import drive
drive.mount('/content/drive')

!wget http://prdownloads.sourceforge.net/ta-lib/ta-lib-0.4.0-src.tar.gz
!tar -xzvf ta-lib-0.4.0-src.tar.gz
!cd ta-lib && ./configure --prefix=/usr && make && make install
!pip install Ta-Lib

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import talib

# 訓練データのファイルパス
path = "***"
# 訓練データの読み込み
df = pd.read_csv(f"{path}/data/train_data.csv")
# indexを"Date"列へ変更
df["Date"] = pd.to_datetime(df["Date"], format="%Y/%m/%d")
df.set_index(keys="Date", inplace=True)

print(f"df.shape: {df.shape}")
df.head(2)

class TechnicalIndicators():
    """
    Create technical indicators by Ta-Lib,
    the open-source library(https://mrjbq7.github.io/ta-lib/index.html).

    Returns:
        dict_df(dict): dictionary of dataframe of each ticker
            key: ticker name
            value: dataset as DataFrame
    
    Indicator:
        MACD
        RSI
    """
    def __init__(self, df, ticker):
        """
        Args:
            df(DataFrame): Dataset of stock prices for each stock
            ticker(list): Ticker, e.g. MSFT, FB..
        """
        self.df = df
        self.ticker = ticker

        # for return
        self.dict_df = {}
    
    def __call__(self):
        """
        Create each indicator
        """
        for ticker in self.ticker:
            # Stock price of the ticker
            price = self.df[ticker]
            # Data of one of the tickers
            self.df_each = pd.DataFrame()
            self.df_each['price'] = price

            # Create indicators
            self.macd(price, fastperiod=12, slowperiod=26, signalperiod=9)
            self.rsi(price, timeperiod=14)

            # Add the created DataFrame to the dictionary
            self.dict_df[ticker] = self.df_each
        
        return self.dict_df
    
    """
    Momentum-indicator Part
    """
    def macd(self, price, fastperiod=12, slowperiod=26, signalperiod=9):
        """
        MACD - Moving Average Convergence/Divergence
            Ref. http://www.tadoc.org/indicator/MACD.htm
        Args:
            price: price of close values
            fastperiod: 
            slowperiod: 
            signalperiod: 
        """
        macd, macdsignal, macdhist = talib.MACD(price, fastperiod=12, slowperiod=26, signalperiod=9)
        self.df_each['macd'] = macd
        self.df_each['macdsignal'] = macdsignal
        self.df_each['macdhist'] = macdhist
    
    def rsi(self, price, timeperiod=14):
        """
        RSI - Relative Strength Index
            Ref. http://www.tadoc.org/indicator/RSI.htm
        Args:
            price: price of close values
        """
        rsi = talib.RSI(price, timeperiod=14)
        self.df_each['rsi'] = rsi
    
    """
    Plot
    """
    def plot_indicator(self, ticker, indicator):
        plt.plot(self.dict_df[ticker][indicator], label=indicator)
        plt.xlabel('date')
        plt.ylabel('indicator value')
        plt.legend()
        plt.show()

    def plot_macd(self, macd, macdsignal, macdhist):
        plt.plot(macd, label='macd')
        plt.plot(macdsignal, label='macdsignal')
        plt.bar(macdhist.index,macdhist, label='macdhist')
        plt.xlabel('date')
        plt.ylabel('indicator value')
        plt.legend()
        plt.show()

# ti: technical
ticker = ['VGSH', 'JEF']
ti = TechnicalIndicators(df, ticker)
df_dict = ti()

# print(df_dict)

# MACD
ticker, indicator = 'JEF', 'macd'
ti.plot_macd(df_dict[ticker]['macd'], df_dict[ticker]['macdsignal'], df_dict[ticker]['macdhist'])

# RSI
ticker, indicator = 'JEF', 'rsi'
ti.plot_indicator(ticker, indicator)