from typing import Any, Dict, Optional

from axelrod.action import Action

from axelrod.evolvable_player import (
    EvolvablePlayer,
    InsufficientParametersError,
    copy_lists,
    crossover_lists,
)

from axelrod.player import Player

C, D = Action.C, Action.D

def is_stochastic_matrix(m, ep=1e-8) -> bool:
    """Checks that the matrix m (a list of lists) is a stochastic matrix."""
    for i in range(len(m)):
        for j in range(len(m[i])):
            if (m[i][j] < 0) or (m[i][j] > 1):
                return False
        s = sum(m[i])
        if abs(1.0 - s) > ep:
            return False
    return True

def normalize_vector(vec):
    s = sum(vec)
    vec = [v / s for v in vec]
    return vec

def mutate_row(row, mutation_probability, rng):
    """, crossover_lists_of_lists
    Given a row of probabilities, randomly change each entry with probability
    `mutation_probability` (a value between 0 and 1).  If changing, then change
    by a value randomly (uniformly) chosen from [-0.25, 0.25] bounded by 0 and
    100%.
    """
    randoms = rng.random(len(row))
    for i in range(len(row)):
        if randoms[i] < mutation_probability:
            ep = rng.uniform(-1, 1) / 4
            row[i] += ep
            if row[i] < 0:
                row[i] = 0
            if row[i] > 1:
                row[i] = 1
    return row

class HMMPlayer(Player):
    """
    Abstract base class for Hidden Markov Model players.

    Names

        - HMM Player: Original name by Marc Harper
    """

    name = "HMM Player"

    classifier = {
        "memory_depth": 1,
        "stochastic": True,
        "long_run_time": False,
        "inspects_source": False,
        "manipulates_source": False,
        "manipulates_state": False,
    }

    def __init__(
        self,
        transitions_C=None,
        transitions_D=None,
        emission_probabilities=None,
        initial_state=0,
        initial_action=C,
    ) -> None:
        Player.__init__(self)
        if not transitions_C:
            transitions_C = [[1]]
            transitions_D = [[1]]
            emission_probabilities = [0.5]  # Not stochastic
            initial_state = 0
        self.initial_state = initial_state
        self.initial_action = initial_action
        self.hmm = SimpleHMM(
            copy_lists(transitions_C),
            copy_lists(transitions_D),
            list(emission_probabilities),
            initial_state,
        )
        assert self.hmm.is_well_formed()
        self.state = self.hmm.state
        self.classifier["stochastic"] = self.is_stochastic()

    def is_stochastic(self) -> bool:
        """Determines if the player is stochastic."""
        # If the transitions matrices and emission_probabilities are all 0 or 1
        # Then the player is stochastic
        values = set(self.hmm.emission_probabilities)
        for m in [self.hmm.transitions_C, self.hmm.transitions_D]:
            for row in m:
                values.update(row)
        if not values.issubset({0, 1}):
            return True
        return False

    def strategy(self, opponent: Player) -> Action:
        """Actual strategy definition that determines player's action."""
        if len(self.history) == 0:
            return self.initial_action
        else:
            action = self.hmm.move(opponent.history[-1])
            # Record the state for testing purposes, this isn't necessary
            # for the strategy to function
            self.state = self.hmm.state
            return action

    def set_seed(self, seed=None):
        super().set_seed(seed=seed)
        # Share RNG with HMM
        # The evolvable version of the class needs to manually share the rng with the HMM
        # after initialization.
        try:
            self.hmm._random = self._random
        except AttributeError:
            pass

class EvolvedHMM5(HMMPlayer):
    """
    An HMM-based player with five hidden states trained with an evolutionary
    algorithm.

    Names:

        - Evolved HMM 5: Original name by Marc Harper
    """

    name = "Evolved HMM 5"

    classifier = {
        "memory_depth": 5,
        "stochastic": True,
        "long_run_time": False,
        "inspects_source": False,
        "manipulates_source": False,
        "manipulates_state": False,
    }

    def __init__(self) -> None:
        initial_state = 3
        initial_action = C
        t_C = [
            [1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0.631, 0, 0, 0.369, 0],
            [0.143, 0.018, 0.118, 0, 0.721],
        ]

        t_D = [
            [0, 1, 0, 0, 0],
            [0, 0.487, 0.513, 0, 0],
            [0, 0, 0, 0.590, 0.410],
            [1, 0, 0, 0, 0],
            [0, 0.287, 0.456, 0.146, 0.111],
        ]

        emissions = [1, 0, 0, 1, 0.111]
        super().__init__(t_C, t_D, emissions, initial_state, initial_action)