from datetime import datetime
from typing import List, Tuple

import numpy as np
import tensorflow.keras.backend as keras_backend

from agent_code.big_bertha_v1.experience_buffer import ExperienceBuffer
from agent_code.big_bertha_v1.features import state_to_features
from agent_code.big_bertha_v1.modifiedtensorboard import ModifiedTensorBoard
from agent_code.big_bertha_v1.parameters import (ACTIONS, BATCH_SIZE,
                                                 EPSILON_DECAY, EPSILON_END,
                                                 EPSILON_START,
                                                 EXPERIENCE_BUFFER_SIZE_MIN,
                                                 GAMMA, REWARDS,
                                                 SAVE_MODEL_EVERY,
                                                 TRAINING_ROUNDS,
                                                 UPDATE_TARGET_MODEL_EVERY,
                                                 UPDATE_TENSORBOARD_EVERY, CURRENT_BUFFER)
from settings import ROWS, COLS


def setup_training(self):
    self.experience_buffer = ExperienceBuffer()  # PrioritizedExperienceBuffer()
    self.experience_buffer.load_samples(CURRENT_BUFFER)

    self.epsilon = EPSILON_START

    self.episode_rewards = []
    self.episode_reward = 0

    self.episode_losses = []
    self.episode_loss = 0

    self.training_timestamp = datetime.now().strftime("%d_%m_%Y_%H_%M_%S")

    self.tensorboard = ModifiedTensorBoard(log_dir="tensorboard_logs/{}".format(self.training_timestamp))


def game_events_occurred(self, old_game_state: dict, self_action: str, new_game_state: dict, events: List[str]):
    if old_game_state is None or new_game_state is None or self_action is None:
        return

    action = ACTIONS.index(self_action)
    reward = reward_from_events(events)
    reward += get_movement_reward(old_game_state['self'][3], new_game_state['self'][3], old_game_state['field'],
                                  old_game_state['bombs'], old_game_state['coins'], old_game_state['others'])
    # td_error = get_td_error(self, old_game_state, action, reward, new_game_state)

    self.experience_buffer.remember(
        state_to_features(old_game_state),
        action,
        reward,
        state_to_features(new_game_state)
    )

    update_q_values(self)
    self.epsilon = self.epsilon * EPSILON_DECAY if self.epsilon > EPSILON_END else EPSILON_END
    self.episode_reward += reward


def end_of_round(self, last_game_state: dict, last_action: str, events: List[str]):
    self.episode_rewards.append(self.episode_reward)
    self.episode_reward = 0

    self.episode_losses.append(self.episode_loss)
    self.episode_loss = 0

    if not last_game_state["round"] % UPDATE_TENSORBOARD_EVERY:
        self.tensorboard.update_stats(
            average_reward=sum(self.episode_rewards[-UPDATE_TENSORBOARD_EVERY:]) / UPDATE_TENSORBOARD_EVERY,
            average_loss=sum(self.episode_losses[-UPDATE_TENSORBOARD_EVERY:]) / UPDATE_TENSORBOARD_EVERY,
            epsilon=self.epsilon,
            learning_rate=keras_backend.eval(self.online_model.optimizer.lr)
        )
    self.tensorboard.step += 1

    if not last_game_state["round"] % UPDATE_TARGET_MODEL_EVERY:
        self.target_model.set_weights(self.online_model.get_weights())

    if not last_game_state["round"] % SAVE_MODEL_EVERY:
        self.online_model.save("models/{}".format(self.training_timestamp))
        print("EXPERT BUFFER SIZE: {} of {}".format(self.experience_buffer.n_entries, self.experience_buffer.size))
        self.experience_buffer.save_samples("buffers/{}".format(self.training_timestamp))

    if last_game_state["round"] == TRAINING_ROUNDS:
        self.online_model.save("models/{}".format(self.training_timestamp))
        print("EXPERT BUFFER SIZE: {} of {}".format(self.experience_buffer.n_entries, self.experience_buffer.size))
        self.experience_buffer.save_samples("buffers/{}".format(self.training_timestamp))


def reward_from_events(occurred_events: List[str]) -> int:
    reward_sum = 0
    for event in occurred_events:
        reward_sum += REWARDS[event]
    return reward_sum


def get_td_error(self, old_game_state: dict, action: int, reward: float, new_game_state: dict) -> float:
    old_game_state = np.array([state_to_features(old_game_state)])
    new_game_state = np.array([state_to_features(new_game_state)])

    q_next_actual = self.target_model.predict(new_game_state)
    q_next_train = self.online_model.predict(new_game_state)
    q_target_old = self.online_model.predict(old_game_state)
    q_target_new = q_target_old.copy()

    q_target_new[0, action] = reward + GAMMA * q_next_actual[0, np.argmax(q_next_train, axis=1)]
    return np.sum(np.abs(q_target_old - q_target_new), axis=1)[0]


def get_movement_reward(old_pos: Tuple[int, int], new_pos: Tuple[int, int], field: np.ndarray,
                        bombs: List[Tuple[Tuple[int, int], int]], coins: List[Tuple[int, int]], others: List[Tuple[str, int, bool, Tuple[int, int]]]) -> float:

    movement_reward = 0

    # TODO: weight with bomb timer
    danger_tile_coordinates = []
    for bomb in bombs:
        bomb_x = bomb[0][0]
        bomb_y = bomb[0][1]
        for x, y in zip([1, 0, -1, 0], [0, 1, 0, -1]):
            for r in range(1, 4):
                if bomb_x + x*r >= 0 and bomb_x + x*r < ROWS and bomb_y + y*r >= 0 and bomb_y + y*r < COLS and field[bomb_x + x*r, bomb_y + y*r] == 0:
                    danger_tile_coordinates.append((bomb_x + x*r, bomb_y + y*r))
        danger_tile_coordinates.append((bomb_x, bomb_y))

    if new_pos in danger_tile_coordinates:
        movement_reward -= 0.3

    if old_pos in danger_tile_coordinates and new_pos not in danger_tile_coordinates:
        movement_reward += 0.3

    old_pos = np.asarray(old_pos)
    new_pos = np.asarray(new_pos)

    # TODO: add enemy position recognition
    if len(coins) > 0:
        coins = np.asarray(coins)
        coin_to_collect = np.argmin(np.sum(np.abs(coins - old_pos), axis=1))

        old_distance = np.sum(np.abs(coins[coin_to_collect] - old_pos))
        new_distance = np.sum(np.abs(coins[coin_to_collect] - new_pos))
        if new_distance < old_distance:
            movement_reward += 0.1
        else:
            movement_reward -= 0.1

    if len(others) > 0:
        others = np.asarray([other[3] for other in others])
        opponent_to_attack = np.argmin(np.sum(np.abs(others - old_pos), axis=1))

        old_distance = np.sum(np.abs(others[opponent_to_attack] - old_pos))
        new_distance = np.sum(np.abs(others[opponent_to_attack] - new_pos))
        if new_distance < old_distance:
            movement_reward += 0.05
        else:
            movement_reward -= 0.05

    return movement_reward


def update_q_values(self):
    if self.experience_buffer.n_entries < EXPERIENCE_BUFFER_SIZE_MIN:
        return

    states, actions, rewards, new_states = self.experience_buffer.sample()  # states, actions, rewards, new_states, weights, indices = self.experience_buffer.sample()

    qs_next_actual = self.target_model.predict(new_states)
    qs_next_train = self.online_model.predict(new_states)
    qs_target_old = self.online_model.predict(states)
    qs_target_new = qs_target_old.copy()

    batch_index = np.arange(BATCH_SIZE)

    qs_target_new[batch_index, actions] = rewards + GAMMA * qs_next_actual[batch_index, np.argmax(qs_next_train, axis=1)]

    # In case of bad performance one could try to recalculate the sumtree every x times
    # First approximations will still be accurate in the first few rounds
    td_errors = np.sum(np.abs(qs_target_new - qs_target_old), axis=1)
    # self.experience_buffer.update(indices, td_errors)
    self.episode_loss += np.sum(td_errors)
    _ = self.online_model.fit(states, qs_target_new, verbose=0)  # _ = self.online_model.train_on_batch(states, qs_target_new, weights)
