# -*- coding: utf-8 -*-
"""Reinforcement_Learning.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1MMpbMt2NUVvRbW2LVB4LZxlEnlfvInht
"""

import math
import random
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt
import pandas_datareader as pdr

from tqdm import tqdm_notebook , tqdm
from collections import deque

tf.__version__

"""# Construct AI"""

class AI_Trader():
  def __init__(self, state_size, action_space=3, model_name="AITrader"):
    self.state_size = state_size
    self.action_space = action_space
    self.memory = deque(maxlen=2000)
    self.model_name= model_name
    self.gamma = 0.95
    self.epsilon = 1
    self.epsilon_final = 0.01
    self.epsilon_decay = 0.995
    self.model = self.modelBuilder()

  def modelBuilder(self):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(units=32,activation='relu',input_dim= self.state_size))
    model.add(tf.keras.layers.Dense(units=64,activation='relu'))
    model.add(tf.keras.layers.Dense(units=128,activation='relu'))
    model.add(tf.keras.layers.Dense(units=self.action_space,activation='linear'))
    model.compile(loss='mse',optimizer = tf.keras.optimizers.Adam(learning_rate=0.001))
    return model
  
  def trade(self,state):
    if random.random() <= self.epsilon:
      return random.randrange(self.action_space)

    actions = self.model.predict(state)
    return np.argmax(actions[0])
  
  def batch_train(self,bath_size):
    batch = []
    for i in range(len(self.memory) - batch_size+1,len(self.memory)):
      batch.append(self.memory[i])

    for state,action,reward,nextState,done in batch:
      if not done:
        reward = reward + self.gamma * np.argmax(self.model.predict(next_state)[0])

      target = self.model.predict(state)
      target[0][action] = reward

      self.model.fit(state, target,epochs=1, verbose=False)
    
    if self.epsilon > self.epsilon_final:
      self.epsilon *= self.epsilon_decay

"""## Pre-Processing DataBa
se

## Sigmoid Function
"""

def sigmoid(x):
  return 1/(1 + math.exp(-x))

sigmoid(0.5)

"""## Price Format"""

def stocks_price_format(n):
  return "- ${0:2f}".format(abs(n)) if n < 0 else "${0:2f}".format(abs(n))

stocks_price_format(100)

"""## Load DataBase"""

dataset = pdr.DataReader('AAPL',data_source='yahoo')

dataset.head()

str(dataset.index[0]).split()[0]

str(dataset.index[-1]).split()[0]

def datasetLoader(stock_name):
  dataset = pdr.DataReader(stock_name,data_source='yahoo')
  start_date = str(dataset.index[0]).split()[0]
  end_date = str(dataset.index[-1]).split()[0]
  close = dataset['Close']
  return close

"""## Ambient States"""

def stateCreator(data,timestep,window_size):
  starting_id = timestep - window_size + 1
  window_data = data[starting_id:timestep+1] if (starting_id >= 0) else  - starting_id * [data[0]] + list(data[0:timestep+1])

  state = []

  for i in range(window_size-1):
    state.append(sigmoid(window_data[i+1]-window_data[i]))
  
  return np.array([state])

"""## Load Database"""

stock_name = 'AAPL'
data = datasetLoader(stock_name=stock_name)

s = stateCreator(data=data,timestep=20,window_size=5)

s

"""# Train

## Hyper-Params
"""

window_size = 10
episodes = 1000
batch_size = 32
data_samples = len(data-1)

trader = AI_Trader(window_size)

trader.model.summary()

for episode in range(1,episodes+1): 
  print(f"Episode: {episode}/{episodes}")
  state = stateCreator(data,0,window_size + 1)
  total_profit = 0
  trader.inventory = []
  for t in tqdm(range(data_samples)):
    action = trader.trade(state=state)
    next_state = stateCreator(data,t+1,window_size+1)
    reward = 0

    if action == 1:
      #Buy Stock
      trader.inventory.append(data[t])
      print(f'AI Trader bought: {stocks_price_format(data[t])}')
    elif action == 2 and len(trader.inventory) > 0:
      buyPrice = trader.inventory.pop(0)
      reward = max(data[t]-buyPrice,0)
      totalProfit = data[t] - buyPrice
      print(f'AI Trader sold: {data[t]} Profit: {stocks_price_format(data[t]-buyPrice)}')
    if t == data_samples - 1:
      done = True
    else: 
      done = False
    
    trader.memory.append((state,action,reward,next_state,done))
    state = next_state

    if done:
      print('#############')
      print('#############')
      print(f'Total Profit{totalProfit}')

    if len(trader.memory) > batch_size:
      trader.batch_train(batch_size)

  if episode % 10 ==0:
    trader.model.save(f'mode_{episode}.h5')