import numpy as np
from numpy.random import uniform
from a_player import APlayer, Reason
from helper import sample, create_empty_state, to_state_string

float_formatter = "{:.3f}".format
np.set_printoptions(formatter={'float_kind': float_formatter})


class MachinePlayer(APlayer):
    def __init__(self, mark='X', with_debug=False, values=None):
        APlayer.__init__(self, mark)
        self.p_exp = 0.2
        self.alpha = 0.1
        self.with_debug = with_debug
        if values is None:
            self.init_values()
        else:
            self.values = values

    def init_values(self):
        self.values = {}

    def print_state_table(self):
        count = 0
        for key in self.values:
            print(f"{self.mark}: {count:05d}: {key} = {self.values[key]:0.3f}")
            count += 1

    @staticmethod
    def to_key(state: np.array):
        key = ''
        for st in state.reshape(state.size):
            key += st
        return key

    def get_value(self, state: np.array):
        key = self.to_key(state)
        try:
            result = self.values[key]
        except KeyError:
            result = 0.5

        return result

    def set_value(self, state: np.array, value):
        self.values[self.to_key(state)] = value

    def reward(self, value):
        self.set_value(self.state, value)

    @staticmethod
    def get_potential_moves(state: np.array) -> np.array:
        st = state.reshape(state.size)
        indices = [idx for idx, s in enumerate(st) if '-' in s]
        return np.array(indices)

    def new_game(self):
        self.state = None
        self.state_last = None

    def move(self, state: np.array):
        do_sample = False
        values = np.array([])
        # get possible move
        moves = self.get_potential_moves(state)
        can_move = moves.size > 0
        best_move = None
        best_value = -1
        if can_move:
            for move in moves:
                # create hypothetical next state
                state_next = self.state_from_move(state.copy(), move)
                # evaluate value
                value = self.get_value(state_next)
                if best_value < value:
                    best_value = value
                    best_move = move
                values = np.append(values, value)

            next_move = best_move
            is_exp = False
            if uniform() < self.p_exp:
                is_exp = True
                index = np.random.randint(len(moves))
                next_move = moves[index]
            elif do_sample:
                index = sample(values)
                next_move = moves[index]

            if self.with_debug:
                print(f"{self.mark}: Values    = {values}")
                print(f"{self.mark}: Moves     = {moves+1}")
                print(f"{self.mark}: Best move = {best_move+1}")
                print(f"{self.mark}: Next move = {next_move+1}, is_exp={is_exp}")

            if self.state is not None:
                self.state_last = self.state.copy()

            self.state = self.state_from_move(state.copy(), next_move)

            # Learn
            if not is_exp and self.state_last is not None:
                v0 = self.get_value(self.state_last)
                v1 = self.get_value(self.state)
                d = v0 + self.alpha*(v1-v0)
                if d > 0:
                    self.set_value(self.state_last, d)
                if self.with_debug:
                    print(f"{self.mark}: Learned {d:0.3f}")

        return self.state, can_move

    def state_from_move(self, state: np.array, field):
        state.reshape(state.size)[field] = self.mark
        return state


class PlayerProvider:
    def __init__(self, p1: APlayer, p2: APlayer):
        self.players = [p1, p2]

    def choose(self) -> list[APlayer]:
        # Wer fängt an?
        players = self.players
        if uniform() < 0.5:
            players.reverse()

        result = players
        players[0].new_game()
        players[1].new_game()

        return result


def play(player_provider: PlayerProvider, k_max=10000, with_print=False):
    for k in range(0, k_max):
        players = player_provider.choose()
        state = create_empty_state()
        move = 1
        run = True
        other_player = players[-1]
        while run:
            for player in players:
                if with_print:
                    print(f"---------------------------------------------------")
                    print(f"- Game {k:06d}, Move {move} -----------------------------")
                    print(f"---------------------------------------------------")
                last_state = state
                state, has_moved = player.move(state)
                if with_print:
                    print(to_state_string(last_state, state))
                if not has_moved:
                    if with_print:
                        print(f"{player.mark}: No more moves")
                    if isinstance(player, MachinePlayer):
                        player.reward(0.0)
                    if isinstance(other_player, MachinePlayer):
                        other_player.reward(0.0)
                    run = False
                if player.has_won(state):
                    if with_print:
                        print(f"{player.mark}: Has won the game")
                    if isinstance(player, MachinePlayer):
                        player.reward(1.0)
                    if isinstance(other_player, MachinePlayer):
                        other_player.reward(0.0)
                    run = False

                other_player = player
                if not run:
                    break

                move += 1


px = MachinePlayer(mark='X', with_debug=False)
po = MachinePlayer(mark='O', with_debug=False)

do_training = 1
if do_training:
    players = PlayerProvider(px, po)
    play(players, 40000, False)

players = PlayerProvider(MachinePlayer(mark='X', with_debug=True, values=px.values), MachinePlayer(mark='O', with_debug=True, values=po.values))
play(players, 1000, True)

px.print_state_table()