import random
import scipy.stats

import tests.common as common

class SequenceGenerator:
    """nb_gm_003:
        - pseudo-random sampling.
        - fixed number of targets.
        - non-skewed stimulus distribution.
        - fixed T:L ratio.

    """

    def __init__(
        self,
        choices,
        n
    ):
        self.trials, self.targets_ratio, self.lures_ratio, self.choices, self.n = None, None, None, choices, n
        self.seq = []

        # create norm distributions for cost functions
        self.skewness_norm = None
        self.lures_norm = None
        self.targets_norm = None

        self.current_trial = None

    def reset(self, trials, targets, lures, **kwargs):
        self.trials = trials
        self.targets_ratio = targets / trials
        self.lures_ratio = lures / trials
        self.n = kwargs.get("n", self.n)

        self.targets_norm = scipy.stats.norm(self.targets_ratio, 1.0)
        self.lures_norm = scipy.stats.norm(self.lures_ratio, 1.0)
        self.skewness_norm = scipy.stats.norm(0, 0.5)

        self.seq = []
        self.current_trial = 0

    def next_trial(self):
        while self.current_trial >= len(self.seq):
            chunk_size = 3 if len(self.seq) + 3 <= self.trials else self.trials-len(self.seq)
            self.seq += self.best_chunk(chunk_size)
        self.current_trial += 1
        print(self.current_trial)
        return self.seq[self.current_trial-1] if self.current_trial<=self.trials else None

    def generate(self, trials, targets, lures):
        self.reset(trials, targets, lures)
        while not self.seq or len(self.seq) < self.trials:
            #self.seq += self.best_choice()
            chunk_size = 3 if len(self.seq) + 3 <= self.trials else self.trials-len(self.seq)
            self.seq += self.best_chunk(chunk_size)
        return self.seq

    def best_chunk(self, chunk_size) -> list:
        from itertools import permutations
        min_cost, best_chunk = None, None
        best_chunk_size = min(len(self.choices), chunk_size)
        chunks = list(permutations(self.choices, best_chunk_size))
        random.shuffle(chunks)
        for chunk in chunks:
            cost = self.cost(self.seq + list(chunk))
            if min_cost is None or cost < min_cost:
                min_cost, best_chunk = cost, chunk
        return list(best_chunk)

    def best_choice(self) -> list:
        best_choice, min_cost = None, None
        random.shuffle(self.choices)  # to avoid ordering effect
        for choice in self.choices:
            cost = self.cost(self.seq + [choice])
            if min_cost is None or cost < min_cost:
                min_cost, best_choice = cost, choice
        return [best_choice]

    def cost(self, seq):
        targets, lures = common.count_targets_and_lures(seq, self.n)
        #print(seq, targets, lures)
        c1, c2, c3 = self.skewness_cost(seq), self.targets_cost(targets/len(seq)), self.lures_cost(lures/len(seq))
        return c1+c2+c3

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

    def targets_cost(self, targets_ratio):
        return 1.0 - self.targets_norm.pdf(targets_ratio) / self.targets_norm.pdf(self.targets_ratio)

    def lures_cost(self, lures_ratio) -> float:
        return 1.0 - self.lures_norm.pdf(lures_ratio) / self.lures_norm.pdf(self.lures_ratio)