import numpy as np
from numpy.random import uniform

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


def softmax(values: np.array, eps=1.0E-6) -> np.array:
    result = (values + eps) / (sum(values + eps))
    return result


def sample(values: np.array, with_debug=False):
    # Normalize and sort
    probs = softmax(values)
    sorted_indices = np.argsort(probs)
    sorted_probs = probs[sorted_indices]
    z = uniform()
    if with_debug:
        print(f"Probs={probs}")
        print(f"Z={z:.3f}")
    index = None
    p_sum = 0
    for idx, p in enumerate(sorted_probs):
        p_sum += p
        if z <= p_sum:
            index = sorted_indices[idx]
            break

    return index


def test_sample(data):
    p = np.array(data)
    c = np.array([0]*len(data))

    for i in range(0, 1000):
        index = sample(p, with_debug=False)
        c[index] += 1

    print(c)


def to_state_string(state, state_nex=None):
    sp = '      '
    sp_arrow = '  =>  '
    sp_ = [sp, sp_arrow, sp]

    def col(str_in, state):
        str_out = str_in
        for c in range(0, 3):
            char = state[3*r+c]
            if char == '-':
                char = ' '
            str_out += "|" + char
        str_out += '|'
        return str_out

    state_str = ''
    for r in range(0, 3):
        state_str = col(state_str, state)

        if state_nex is not None:
            state_str += sp_[r]
            state_str = col(state_str, state_nex)

        if r != 2:
            state_str += '\x0A'

    return state_str


class APlayer(object):
    def __init__(self, mark):
        self.mark = mark
        self.state = None
        self.state_opp = None
        self.state_last = None

    def move(self, state):
        return state, False

    def has_won(self, state):
        ref = self.mark + self.mark + self.mark

        substr = state[0:3]
        if substr in ref:
            return True

        substr = state[3:6]
        if substr in ref:
            return True

        substr = state[6:9]
        if substr in ref:
            return True

        substr = state[0:9:3]
        if substr in ref:
            return True

        substr = state[1:9:3]
        if substr in ref:
            return True

        substr = state[2:9:3]
        if substr in ref:
            return True

        substr = state[0:9:4]
        if substr in ref:
            return True

        substr = state[6:0:-2]
        if substr in ref:
            return True

        return False


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.9
        self.with_debug = with_debug
        if values is None:
            self.values = {}
        else:
            self.values = values

    def get_value(self, state):
        try:
            result = self.values[state]
        except KeyError:
            result = 0

        return result

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

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

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

    def move(self, state):
        values = np.array([])
        # get possible move
        moves = self.get_potential_moves(state)
        can_move = moves.size > 0
        self.state_opp = state
        if can_move:
            for move in moves:
                # create hypothetical next state
                state_next = self.state_from_move(state, move)
                # evaluate value
                value = self.get_value(state_next)
                values = np.append(values, value)

            is_exp = False
            if uniform() < self.p_exp:
                is_exp = True
                index = np.random.randint(len(moves))
            else:
                index = sample(values)

            move = moves[index]
            if self.with_debug:
                print(f"{self.mark}: Values = {values}")
                print(f"{self.mark}: Moves = {moves+1}")
                print(f"{self.mark}: Move = {move+1}, is_exp={is_exp}")

            self.state_last = self.state
            self.state = self.state_from_move(state, move)

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

        return self.state, can_move

    def state_from_move(self, state, field):
        return state[:field] + self.mark + state[field + 1:]


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


def choose_player(player_provider: PlayerProvider) -> list[APlayer]:
    # Wer fängt an?
    p1 = player_provider.p1
    p2 = player_provider.p2
    if uniform() < 0.5:
        result = [p1, p2]
    else:
        result = [p2, p1]

    return result


def play(player_provider: PlayerProvider, k_max=10000, with_print=False):
    for k in range(0, k_max):
        players = choose_player(player_provider)
        state = "---------"
        move = 1
        run = True
        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")
                    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)
                    run = False

                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, 20000, False)

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

print("Testing")
test_sample([0.7, 0.1, 0.1, 0.1])
test_sample([0.1, 0.1, 0.2, 0.1])
test_sample([0.2, 0.8])
test_sample([0.1, 0.1])
test_sample([0.3, 0.0, 0.2, 0.0])
test_sample([0.0, 0.0, 0.0, 0.0])
test_sample([0.5, 0.0, 0.2, 0.3])
test_sample([0.3, 0.1, 0.4, 0.2])