import random
import scipy.stats


class SequenceGenerator:
    """Generate a sequence progressively according to a predefined TL ratio and an even distribution"""

    def __init__(
        self,
        choices,
        trials,
        desired_tl=4.0,
        n=3,
        targets_ratio=0.33
    ):
        self.desired_tl, self.choices, self.n, self.targets_ratio = desired_tl, choices, n, targets_ratio
        self.sequence = list()
        self.norm_even_dist = scipy.stats.norm(0, trials/2)
        self.norm_targets_ratio_dist = scipy.stats.norm(targets_ratio, 0.5)
        self.norm_tl_ratio_dist = scipy.stats.norm(desired_tl, trials/2)
        self.trials = None

    def generate(self, trials):
        self.trials = trials
        while not self.sequence or len(self.sequence) < self.trials:
            self.sequence = self.sequence + self.__find_best_choice(self.sequence, self.choices)
        return self.sequence

    def next_trial(self):
        if self.sequence and len(self.sequence) >= self.trials:
            return None
        self.sequence = self.sequence + self.__find_best_choice(self.sequence, self.choices)
        return self.sequence[-1]

    def __find_best_choice(self, seq: list, choices: list) -> list:
        import sys
        min_cost, best_choice = sys.float_info.max, None
        random.shuffle(choices)  # to avoid ordering effect
        for choice in choices:
            cost = self.cost(seq + choice)
            if cost < min_cost:
                min_cost = cost
                best_choice = choice
        return best_choice

    def even_distribution_cost(self, seq):
        even_ratio = self.trials / len(self.choices)
        costs = {c: abs(seq.count(c) - even_ratio)/self.trials for c in self.choices}
        max_dist = max(list(costs.values()))
        return 1.0 - self.norm_even_dist.pdf(max_dist)

    def cost(self, seq):
        """
        Calculate overall fitness of a sequence (block of trials).
        Right now it's a cost function, so we try to minimize this cost.
        :param seq:
        :return:
        """

        targets, lures = self.count_targets_and_lures(seq)
        targets_cost = 1.0 - self.norm_targets_ratio_dist.pdf(targets/self.trials)
        tl_cost = 1.0 - self.norm_tl_ratio_dist.pdf(self.tl(seq))
        even_cost = self.even_distribution_cost(seq)
        return targets_cost + tl_cost + even_cost

    def count_targets_and_lures(self, seq):
        n = self.n
        targets = 0.0
        lures = 0.0
        for index in range(n, len(seq)):
            if seq[index] == seq[index - n]:
                targets += 1.0
            elif seq[index] == seq[index - (n-1)] or seq[index] == seq[index - (n+1)]:
                lures += 1.0
        return targets, lures

    def tl(self, seq):
        """Calculates the T/L ratio in a block of trials."""
        targets, lures = self.count_targets_and_lures(seq)
        if lures < 0.01:  # avoid division by zero
            lures = 0.01
        return targets/lures