# S -- A, R --> S'
import gym
import numpy as np
import pandas as pd
import math


## discrete - table
class Table(object):
    def __init__(self, actions, learning_rate=0.01, reward_decay=0.9, e_greedy=0.9):
        self.actions = actions  # a list
        self.lr = learning_rate
        self.gamma = reward_decay
        self.epsilon = e_greedy
        self.q_table = pd.DataFrame(columns=self.actions, dtype=np.float64)
        self.lastAction = actions[0]
        self.lastState = 0
        self.check_state_exist(self.lastState)
    
    def ob2state(self,ob):
        # for moutain car: discrete
        # position -1.2 0.6
        # speed 0.07
        # state = str(math.ceil((ob[0]+1.2)/0.01)) +" "+ str(math.ceil((ob[1]+0.07)/0.01))
        # for fish
        state = str(math.ceil(ob[0]/0.01)) +" "+ str(math.ceil(ob[1]/0.01))+" "+ str(math.ceil(ob[2]/0.01))+" "+ str(math.ceil(ob[3]/0.01))
        return state 

    def act(self, ob, reward, done): 
        # last action -> observe r,s'
        # Q(s,a) <- Q(s,a) + alpha*(r + gamma*maxQ(s',a') - Q(s,a))
        s = self.lastState
        a = self.lastAction
        s_ = self.ob2state(ob)
        self.check_state_exist(s_)
        self.lastState = s_

        action = self.choose_action(s_)
        self.learn(s,a,s_,reward,action)

        self.lastAction = action
        return action
    
    def learn(self, s, a, s_, r):
        pass

    def choose_action(self,s_):
        # Choose a from s using policy derived from Q (epsilon-greedy)
        if np.random.uniform() < self.epsilon:
            # choose best action
            state_action = self.q_table.loc[s_, :]
            # some actions may have the same value, randomly choose on in these actions
            action = np.random.choice(state_action[state_action == np.max(state_action)].index)
        else:
            # choose random action
            action = np.random.choice(self.actions)
        self.lastAction = action
        return action

    def check_state_exist(self, state):
        if state not in self.q_table.index:
            # append new state to q table
            self.q_table = self.q_table.append(
                pd.Series(
                    [0]*len(self.actions),
                    index=self.q_table.columns,
                    name=state,
                )
            )
    

# Q-learning
class QLearningTable(Table):
    def learn(self, s, a, s_, r,_):
        self.q_table.loc[s,a] += self.lr*(reward + self.gamma*self.q_table.loc[s_, :].max() - self.q_table.loc[s,a])

# SARSA
class SARSA(Table):
    def learn(self, s, a, s_, r, a_):
        self.q_table.loc[s,a] += self.lr*(reward + self.gamma*self.q_table.loc[s_, a_] - self.q_table.loc[s,a])

## continouos State
import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self, n_states, n_actions):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(n_states, 8)
        self.out = nn.Linear(8, n_actions)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        actions_value = self.out(x)
        return actions_value

class DQN(object):
    def __init__(self, env, epsilon = 0.9, Cmemo=8000, gamma=0.9, lr=0.01, bz=128, ftar=100):
        self.n_states = env.observation_space.shape[0]
        self.n_actions = env.action_space.n

        self.eval_net = Net(self.n_states, self.n_actions)   # 实时更新预测
        self.target_net = Net(self.n_states, self.n_actions) # “现实”值锁定
        self.target_net_replace_iter = ftar
        self.learn_step_counter = 0     # 用于 target 更新计时

        self.epsilon = epsilon
        self.gamma = gamma

        self.lr = lr
        self.batch_zise = bz
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=self.lr)    # torch 的优化器
        self.loss_func = nn.MSELoss()   # 误差公式

        self.memory_capacity = Cmemo # 记忆库大小
        self.memory_counter = 0         # 记忆库记数
        
        self.lastAction = 1
        self.lastState = None

        self.ep_loss = 0

        self.use_ex = False
        if self.use_ex:
            self.memory = np.load('memo.npy')
        else:
            self.memory = np.zeros((self.memory_capacity, self.n_states * 2 + 2))     # 初始化记忆库
    
    def act(self, s, r, done, learning = True):
        if not self.use_ex:
            self.store_transition(self.lastState, self.lastAction, r, s)
        if learning and(self.use_ex or self.memory_counter > self.memory_capacity):
            self.learn() # 记忆库满了就进行学习
        a = self.choose_action(s)
        self.lastAction = a
        self.lastState = s
        return a

    def choose_action(self, x):
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        # 这里只输入一个 sample
        if np.random.uniform() < self.epsilon:   # 选最优动作
            actions_value = self.eval_net.forward(x)
            action = torch.max(actions_value, 1)[1].data.numpy()[0]#[0, 0]     # return the argmax
        else:   # 选随机动作
            action = np.random.randint(0, self.n_actions)
        return action

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, [a, r], s_))
        # 如果记忆库满了, 就覆盖老数据
        index = self.memory_counter % self.memory_capacity
        self.memory[index, :] = transition
        self.memory_counter += 1

    def learn(self):
        # target net 参数更新
        if self.learn_step_counter % self.target_net_replace_iter == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1

        # 抽取记忆库中的批数据
        sample_index = np.random.choice(self.memory_capacity, self.batch_zise)
        b_memory = self.memory[sample_index, :] 
        b_s = torch.FloatTensor(b_memory[:, :self.n_states])
        b_a = torch.LongTensor(b_memory[:, self.n_states:self.n_states+1].astype(int))
        b_r = torch.FloatTensor(b_memory[:, self.n_states+1:self.n_states+2])
        b_s_ = torch.FloatTensor(b_memory[:, -self.n_states:])

        # 针对做过的动作b_a, 来选 q_eval 的值, (q_eval 原本有所有动作的值)
        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)
        q_next = self.target_net(b_s_).detach()     # q_next 不进行反向传递误差, 所以 detach
        q_target = b_r + self.gamma * q_next.max(1)[0].view(self.batch_zise, 1)   # shape (batch, 1)
        loss = self.loss_func(q_eval, q_target)
        self.ep_loss += (float)(loss)

        # 计算, 更新 eval net
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    
    def save_model(self, path):
        torch.save(self.eval_net.state_dict(), path+"eval.pt")
        torch.save(self.target_net.state_dict(), path+"tar.pt")
    
    def load_model(self,path):
        self.eval_net.load_state_dict(torch.load(path+"eval.pt"))
        self.target_net.load_state_dict(torch.load(path+"tar.pt"))


## main ##
#init
done = False
reward = 0
from fish import FishingEnv
env = FishingEnv()
env.fish_max_acc *= 2
#env = gym.make('MountainCar-v0')

# <-----method----->
#agent = SARSA(actions=list(range(env.action_space.n)))
#agent = QLearningTable(actions=list(range(env.action_space.n)))
agent = DQN(env)
# <---------->
#agent.use_ex = True
#agent.load_model("")
total_suc_time = 0
losses = []
i = 0
running = True
episode_count = 10
while running:
    step_used = 0
    agent.ep_loss = 0.0
    ob = env.reset()
    agent.lastState = ob
    Rt = 0
    done = False
    while True:
        if i > episode_count-10:
            env.render()
        step_used += 1
        #reward = -(math.e ** (abs(ob[0]-ob[2])-1))
        action = agent.act(ob, reward, done,False) #False
        if done:
            if env.catch:
                total_suc_time += 1
            break
        ob, reward, done, _ = env.step(action)
        Rt += reward
    if agent.memory_counter >= agent.memory_capacity:
        i += 1
        losses.append(agent.ep_loss/(float)(step_used))
    if i == episode_count:
        running = False
    print("ep"+str(i)+"/"+str(total_suc_time)+", "+str(round(agent.ep_loss/(float)(step_used),3))+", "+str(step_used)+", "+str(Rt))
    #print("ep"+str(i)+"/"+str(total_suc_time)+", "+str(step_used)+", "+str(Rt))

env.close()
#agent.save_model("")

import matplotlib.pyplot as plt
plt.plot(losses)
plt.show()