import numpy as np
from numpy.random import uniform


def softmax(values: np.array) -> np.array:
    result = values / sum(values)
    result = np.sort(result)
    return result


def sample(values: np.array):
    # Normalize and sort
    probs = softmax(values + 1.0E-6)
    z = uniform()
    print(f"Probs={probs}")
    print(f"Z={z}")
    result = None
    p_sum = 0
    for idx, p in enumerate(probs):
        p_sum += p
        if z <= p_sum:
            result = idx
            break

    is_exploration = result != (len(probs) - 1)
    return result, is_exploration


def test_sample():
    p = np.array([0.1, 0.1, 0.3, 0.5])
    c = np.array([0, 0, 0, 0])

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

    print(c)


def to_state_string(state):
    state_str = ''
    for i in range(0, 3):
        for j in range(0, 3):
            char = state[3*i+j]
            if char == '-':
                char = ' '
            state_str += "|"+char
        state_str += '|\x0A'

    return state_str


class Player(object):
    def __init__(self, mark='X'):
        self.values = {}
        self.mark = mark
        self.state_last = None

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

        return result

    def set_value(self, value, state=None):
        if state is None:
            if self.state_last is not None:
                self.values[self.state_last] = value
        else:
            self.values[state] = 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
        state_next = state

        if can_move:
            for field in moves:
                # crate hypothetical next state
                state_next = self.state_from_move(state, field)
                # evaluate value
                value = self.get_value(state_next)
                values = np.append(values, value)

            index, is_exp = sample(values)
            field = moves[index]
            print(f"{player.mark}: Chose {index}")
            state_next = self.state_from_move(state, field)

            # Learn
            if not is_exp and self.state_last is not None:
                v0 = self.get_value(self.state_last)
                v1 = self.get_value(state_next)
                d = max(0, v1-v0)
                if d > 0:
                    self.set_value(0.1*d)
                print(f"{player.mark}: Learned {d}")

            self.state_last = state_next
        return state_next, can_move

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

    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


p_x = Player(mark='X')
p_o = Player(mark='O')


for k in range(0, 10000):
    # Wer fängt an?
    if uniform() < 0.5:
        players = [p_x, p_o]
    else:
        players = [p_o, p_x]
    state = "---------"
    move = 1
    run = True
    last_state = None
    while run:
        last_state = state
        for player in players:
            print(to_state_string(state))
            state, has_moved = player.move(state)

            if not has_moved:
                print(f"{player.mark}: No more moves")
                run = False
            if player.has_won(state):
                print(f"{player.mark}: Has won the game")
                player.set_value(1.0)
                print(to_state_string(state))
                run = False

            if not run:
                break

            move += 1
    for player in players:
        print(f"{player.mark}: Values: {player.values}")