from ple.games.snake import Snake
from ple import PLE
import numpy as np
from agent import Agent
import pygame
import sys


def get_dist(head_x, head_y, obs_x, obs_y):
    return ((head_x - obs_x) ** 2 + (head_y - obs_y) ** 2) ** 0.5


def get_state(state):
    head_x, head_y = state[0], state[1]
    min_dist_walls = min(get_dist(head_x, head_y, head_x, 0), get_dist(head_x, head_y, 0, head_y),
                         get_dist(head_x, head_y, 600, head_y), get_dist(head_x, head_y, head_x, 600))
    return [state[0], state[1], state[2], state[3], min(min(state[4][4:]), min_dist_walls)]


def vision(state):
    my_vision = [[0,0] for _ in range(4)]
    head_x, head_y = state[0], state[1]
    food_x, food_y = state[2], state[3]

    # food
    dist_x, dist_y = head_x - food_x, head_y - food_y
    if abs(dist_y) < 500:
        if dist_x < 0:
            my_vision[3][0] = 10
        else:
            my_vision[2][0] = 10
    if abs(dist_x) < 500:
        if dist_y < 0:
            my_vision[1][0] = 10
        else:
            my_vision[0][0] = 10

    # wall
    if head_x <= 50:
        my_vision[2][1] = -100
    elif 600 - head_x <= 50:
        my_vision[3][1] = -100
    if head_y <= 50:
        my_vision[0][1] = -100
    elif 600 - head_y <= 50:
        my_vision[1][1] = -100

    # body
    for body_x, body_y in state[5][3:]:
        # print(body_x,body_y)
        dist_x = head_x - body_x
        dist_y = head_y - body_y
        if abs(dist_x) <= 50:
            if dist_x > 0:
                my_vision[2][1] = -100
            else:
                my_vision[3][1] = -100
        if abs(dist_y) <= 50:
            if dist_y < 0:
                my_vision[1][1] = -100
            else:
                my_vision[0][1] = -100
    output = []
    [output.extend(item) for item in my_vision]
    output.extend([head_x, head_y, food_x, food_y])
    return output


def prepare_corect_directions(direction):
    direction = str(direction)
    if direction == "Left":
        return {119: "Up", 115: "Down", 97: "Left"}
    if direction == "Right":
        return {115: "Down", 119: "Up", 100: "Right"}
    if direction == "Up":
        return {100: "Right", 97: "Left", 119: "Up"}
    if direction == "Down":
        return {97: "Left", 100: "Right", 115: "Down"}


def process_state(state):
    return np.array([state.values()])


def run():
    game = Snake(600, 600)
    p = PLE(game, state_preprocessor=process_state,force_fps=True, display_screen=True,frame_skip=5)
    print(sys.argv[1])
    agent = Agent(alpha=float(sys.argv[1]), gamma=float(sys.argv[2]), n_actions=3, epsilon=0.99, batch_size=100,
                  input_shape=12, epsilon_dec=0.99999,
                  epsilon_end=0.1,
                  memory_size=50000000, file_name=sys.argv[3], activations=[str(sys.argv[4]), str(sys.argv[5])])
    p.init()
    # agent.load_game()

    scores = []

    for _ in range(100000):
        if p.game_over():
            p.reset_game()
        score = 0
        # state = p.getGameState()
        initial_direction = "Right"
        game_state = np.array(vision(list(p.getGameState()[0])))
        # print(game_state)
        prec_dist = get_dist(game_state[0], game_state[1], game_state[2], game_state[3])

        while not p.game_over():
            old_state = np.array(vision(list(p.getGameState()[0])))

            action = agent.choose_action(old_state)

            possible_directions = prepare_corect_directions(initial_direction)
            possible_directions_tuples = list(zip(possible_directions.keys(), possible_directions.values()))
            direction = possible_directions_tuples[action]
            initial_direction = direction[1]

            reward = p.act(direction[0])
            if reward == -0.1:
                game_state = np.array(vision(list(p.getGameState()[0])))
                curr_dist = get_dist(game_state[0], game_state[1], game_state[2], game_state[3])
                if prec_dist > curr_dist: reward = 1.5
                prec_dist=curr_dist

            print(reward)

            new_state = np.array(vision(list(p.getGameState()[0])))
            agent.add_experience(old_state, action, reward, new_state)
            agent.learn()
            score = p.score()
        scores.append(score)
        print(f"Score for model iteration number _ {str(sys.argv[3])} with learning_rate {sys.argv[1]}, gama {sys.argv[2]}, activations: {sys.argv[4], sys.argv[5]} is score {score}. Epsilon is {agent.epsilon}")
        agent.save_game()


#

if __name__ == '__main__':
    run()