import numpy as np
import pandas as pd
import time
from collections import defaultdict
from argparse import ArgumentParser


class GridWorld:
    def __init__(self, path):
        self.path = path  # loaded file path
        self.grid = np.loadtxt(self.path, dtype='str')  # initial grid
        self.n_row, self.n_col = np.shape(self.grid)  # # of row, # of colum
        self.start_state = list()  # start sate
        self.barrier_list = list()  # list of barrier state
        self.end_state = list()  # lsit of end state
        self.possible_state = list()  # list of reachable states
        self.act_n = ['up', 'down', 'right', 'left']  # name list of possible actions
        self.sta_n = 0  # number of possible states
        self._Q = None  # initial Q table
        self.view_dic = defaultdict(int)  # dictionary of number of reach times in each square during training

    def init_func(self):
        for i in range(self.n_row):
            for j in range(self.n_col):
                value = self.grid[i][j]
                if value == 'X':
                    self.barrier_list.append((i, j))
                else:
                    if value == 'S':
                        self.start_state.append((i, j))
                        self.possible_state.append((i, j))
                    elif value == '0':
                        self.possible_state.append((i, j))
                    else:
                        self.end_state.append((i, j))
        self.sta_n = len(self.possible_state)
        Q_table = np.zeros((self.sta_n, len(self.act_n)))
        self._Q = pd.DataFrame(Q_table, index=self.possible_state, columns=self.act_n, dtype='float')
        for i in self.possible_state:
            self.view_dic[i] = 0


class RLearningAgent(GridWorld):
    def __init__(self, path, learning_rate, gamma, suc_p, reward, time, epsilon=0.9):
        super(RLearningAgent, self).__init__(path)
        self._lr = learning_rate  # how fast the algorithm to converge
        self._gamma = gamma  # discounting rate
        self._sp = suc_p  # action success probability
        self.reward = reward  # reward for each action takes
        self._rt = time  # trainning time
        self._ep = epsilon  # epsilon value, probability of choose max Q in action determination
        self.score = 0  # score each trail get
        self.total_score = 0  # total score of training
        self.n_round = 0  # number of training trails

    def get_max_Q_direction(self, state):
        data = self._Q.loc[[state], :]
        max_value = data.max(axis=1)[0]
        result = data.isin([max_value])
        col_list = list(result.any()[result.any() == True].index)
        index = np.random.randint(0, len(col_list))
        direction = col_list[index]
        return direction

    def determineAction(self, state):
        rand_num = np.random.uniform(0, 1)
        if rand_num < self._ep:
            direction = self.get_max_Q_direction(state)  # direction of highest Q(s,a) value
        else:
            index = np.random.randint(0, len(self.act_n))
            direction = self.act_n[index]  # direction of highest Q(s,a) value
        return direction

    def takeAction_p2(self, state, direction):
        rand_num = np.random.uniform(0, 1)
        if rand_num <= self._sp:  # successful, forward 1 squre
            steps = 1
        elif rand_num >= (1 - (1 - self._sp) / 2):  # random action 2, backward 1 squre
            if direction == 'up':
                direction = 'down'
                steps = 1
            elif direction == 'down':
                direction = 'up'
                steps = 1
            elif direction == 'left':
                direction = 'right'
                steps = 1
            else:
                direction = 'left'
                steps = 1
        else:  # random action 1, jump 2 squres
            steps = 2
        row = state[0]
        col = state[1]
        if steps == 1:
            if direction == 'up':
                next_row = row - steps
                next_col = col
                next_state = (next_row, next_col)
                if (next_row, next_col) in self.barrier_list or next_row < 0:  # reach barrier or boundary
                    next_state = (row, col)  # stay
            elif direction == 'down':
                next_row = row + steps
                next_col = col
                next_state = (next_row, next_col)
                if (next_row, next_col) in self.barrier_list or next_row >= self.n_row:  # reach barrier
                    next_state = (row, col)  # stay
            elif direction == 'right':
                next_row = row
                next_col = col + steps
                next_state = (next_row, next_col)
                if (next_row, next_col) in self.barrier_list or next_col >= self.n_col:  # reach barrier
                    next_state = (row, col)  # stay
            else:
                next_row = row
                next_col = col - steps
                next_state = (next_row, next_col)
                if (next_row, next_col) in self.barrier_list or next_col < 0:  # reach barrier
                    next_state = (row, col)  # stay
        else:
            if direction == 'up':
                next_row = row - steps
                next_col = col
                next_state = (next_row, next_col)
                if (next_row + 1, next_col) in self.barrier_list or next_row + 1 < 0:  # reach barrier or boundary
                    next_state = (row, col)  # stay
                elif (next_row, next_col) in self.barrier_list or next_row < 0:  # 1 squre from barrier or boundary
                    next_state = (next_row + 1, col)  # move up 1 square
            elif direction == 'down':
                next_row = row + steps
                next_col = col
                next_state = (next_row, next_col)
                if (
                next_row - 1, next_col) in self.barrier_list or next_row - 1 >= self.n_row:  # reach barrier or boundary
                    next_state = (row, col)  # stay
                elif (
                next_row, next_col) in self.barrier_list or next_row >= self.n_row:  # 1 squre from barrier or boundary
                    next_state = (next_row - 1, next_col)
            elif direction == 'right':
                next_row = row
                next_col = col + steps
                next_state = (next_row, next_col)
                if (
                next_row, next_col - 1) in self.barrier_list or next_col - 1 >= self.n_col:  # reach barrier or boundary
                    next_state = (row, col)  # stay
                elif (
                next_row, next_col) in self.barrier_list or next_col >= self.n_col:  # 1 squre from barrier or boundary
                    next_state = (next_row, next_col - 1)
            else:
                next_row = row
                next_col = col - steps
                next_state = (next_row, next_col)
                if (next_row, next_col + 1) in self.barrier_list or next_col + 1 < 0:  # reach barrier or boundary
                    next_state = (row, col)  # stay
                elif (next_row, next_col) in self.barrier_list or next_col < 0:  # 1 squre from barrier or boundary
                    next_state = (next_row, next_col + 1)
        return next_state

    def notTerminal(self, state):
        if state in self.end_state:
            return False
        else:
            return True

    def QL_update(self, state, direction, next_state):
        predict_Q = self._Q.at[state, direction]
        if self.notTerminal(next_state):
            next_direction = self.get_max_Q_direction(next_state)  # direction of highest Q(s,a) value in next state
            target_Q = self.reward + self._gamma * self._Q.at[next_state, next_direction]
            self.score += self.reward
            self.view_dic[next_state] += 1
        else:
            target_Q = int(self.grid[next_state[0]][next_state[1]])
            self.score += int(self.grid[next_state[0]][next_state[1]])
        self._Q.at[state, direction] += self._lr * (target_Q - predict_Q)

    def QL_learning_p2(self):
        st = time.time()
        timeRemains = 0
        while timeRemains < self._rt:
            self.score = 0
            self.n_round += 1
            index = np.random.randint(0, len(self.start_state))
            state = self.start_state[index]
            self.view_dic[state] += 1
            while self.notTerminal(state):
                direction = self.determineAction(state)
                next_state = self.takeAction_p2(state, direction)
                self.QL_update(state, direction, next_state)
                state = next_state
            # print('the score of {}th round ------ {}'.format(self.n_round, self.score))  # can be deleted
            et = time.time()
            timeRemains = et - st
            self.total_score += self.score
        return self.total_score / self.n_round

    def SARSA_update(self, state, direction, next_state, next_direction):
        predict_Q = self._Q.at[state, direction]
        if self.notTerminal(next_state):
            target_Q = self.reward + self._gamma * self._Q.at[next_state, next_direction]
            self.score += self.reward
            self.view_dic[next_state] += 1
        else:
            target_Q = int(self.grid[next_state[0]][next_state[1]])
            self.score += int(self.grid[next_state[0]][next_state[1]])
        self._Q.at[state, direction] += self._lr * (target_Q - predict_Q)

    def SARSA_learning(self):
        st = time.time()
        timeRemains = 0
        while timeRemains < self._rt:
            self.score = 0
            self.n_round += 1
            index = np.random.randint(0, len(self.start_state))
            state = self.start_state[index]
            self.view_dic[state] += 1
            while self.notTerminal(state):
                direction = self.determineAction(state)
                next_state = self.takeAction_p2(state, direction)
                if self.notTerminal(next_state):
                    next_direction = self.determineAction(next_state)
                else:
                    next_direction = ''
                self.SARSA_update(state, direction, next_state, next_direction)
                state = next_state
                direction = next_direction
            # print('the score of {}th round ------ {}'.format(self.n_round, self.score)) # can be deleted
            et = time.time()
            timeRemains = et - st
            self.total_score += self.score
        return self.total_score / self.n_round

    def draw_heat_grid_Q(self):
        index_list = self._Q.index
        heat_grid = np.zeros((self.n_row, self.n_col))
        d = pd.DataFrame(0, index=index_list, columns=['Q_max'], dtype='float')
        ter_sum = 0
        for i in index_list:
            max_value = self._Q.loc[[i], :].max(axis=1)[0]
            d.loc[[i], 'Q_max'] = abs(max_value)
        max_Q = d.max()[0]
        d = d / max_Q
        total_Q = d.sum()[0]
        for i in self.end_state:
            row = i[0]
            col = i[1]
            value = int(self.grid[row][col])
            ter_sum += value
            heat_grid[row][col] = value
        for i in self.barrier_list:
            heat_grid[i[0]][i[1]] = float('-inf')
        for i in index_list:
            row = i[0]
            col = i[1]
            heat_grid[row][col] = d.loc[[i], 'Q_max'] / total_Q * (100 - ter_sum)
        return heat_grid

    def draw_heat_grid_V(self):
        index_list = list(self.view_dic.keys())
        heat_grid = np.zeros((self.n_row, self.n_col))
        view_pd = pd.DataFrame.from_dict(self.view_dic, orient='index', columns=['num'])
        total_view = view_pd.sum()[0]
        ter_sum = 0
        for i in self.end_state:
            row = i[0]
            col = i[1]
            value = int(self.grid[row][col])
            ter_sum += value
            heat_grid[row][col] = value
        for i in self.barrier_list:
            heat_grid[i[0]][i[1]] = float('-inf')
        for i in index_list:
            row = i[0]
            col = i[1]
            value = self.view_dic[i] / total_view * (100 - ter_sum)
            heat_grid[row][col] = value
        return heat_grid

    def draw_policy_grid(self):
        index_list = self._Q.index
        policy_grid = self.grid[:]
        for i in index_list:
            row = i[0]
            col = i[1]
            direction = self._Q.loc[[i], :].idxmax(axis=1)[0]
            if direction == 'right':
                policy_grid[row][col] = '>'
            elif direction == 'left':
                policy_grid[row][col] = '<'
            elif direction == 'up':
                policy_grid[row][col] = '^'
            else:
                policy_grid[row][col] = 'v'
        return policy_grid

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-f", dest="file_path", help="the file path for the grid with postfix",
                        default='grid_sample.txt')
    parser.add_argument("-lr", dest="learning_rate", help="learning rate of RL", type=float,
                        default=0.1)
    parser.add_argument("-g", dest="gamma", help="discounting rate", type=float,
                        default=0.9)
    parser.add_argument("-s", dest="action_success", help="action success probability", type=float,
                        default=1.0)
    parser.add_argument("-r", dest="reward", help="RL rewards", type=float,
                        default=-0.04)
    parser.add_argument("-t", dest="run_time", help="training time", type=float,
                        default=10.0)
    parser.add_argument("-e", dest="epsilon", help="epsilon greedy exploration rate", type=float,
                        default=0.9)


    parsed_args = parser.parse_args()


    path = parsed_args.file_path # grid file path
    lr = parsed_args.learning_rate # learning rate
    dr = parsed_args.gamma  # discounting rate
    sr = parsed_args.action_success # action success probability
    rw = parsed_args.reward # reward for each action takes
    run_t = parsed_args.run_time # training time
    epsilon = parsed_args.epsilon # epsilon value, probability of choose max Q in action determination


    # temp = GridWorld(path)
    # temp.init_func()
    # test0 = temp._Q
    # print(test0)

    start_time = time.time()

    LearningClass = RLearningAgent(path, lr, dr, sr, rw, run_t, epsilon)
    LearningClass.init_func()
    original_mtx = LearningClass.grid
    print('The original board is')
    print(original_mtx)

    mean_rw = LearningClass.QL_learning_p2()
    updated_mtx = LearningClass._Q

    heat_mtx = LearningClass.draw_heat_grid_V()
    policy_mtx = LearningClass.draw_policy_grid()

    # print('The total run time is %s' %(time.time()-start_time))



    # print('The updated Q table is:')
    # print(updated_mtx)

    print('The heat map after training is:')
    print(heat_mtx)
    print('The policy after training is:')
    print(policy_mtx)

    print("------------------------------------------")
    # print("The total reward is: %s" % (q_temp.total_score))
    print("The number of times for reaching a terminal is: %s" % (LearningClass.n_round))
    print("The mean reward per trail is: %s" % (mean_rw))