import random
import scipy.stats

import benchmarks.common as common

class SequenceGenerator:
    """nb_gm_003:
        - pseudo-random sampling.
        - specific number of matching trials.
        - even distribution of stimuli.
    """

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

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

    def reset(self, trials, targets):
        self.trials = trials
        self.targets = targets

        self.targets_norm = scipy.stats.norm(targets, 0.5)
        self.skewness_norm = scipy.stats.norm(0, 0.5)

        if self.seq:
            self.seq.clear()

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

    def best_chunk(self, chunk_size=3) -> list:
        from itertools import permutations
        min_cost, best_chunk = None, None
        chunks = list(permutations(self.choices, 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):
        # DEBUG print(self.matchratio_cost(seq), self.evendist_cost(seq))
        return self.skewness_cost(seq) + self.targets_ratio_cost(seq)

    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()))
        cost = 1.0 - (self.skewness_norm.pdf(max_deviation_from_even_dist) / self.skewness_norm.pdf(0))
        return cost

    def targets_ratio_cost(self, seq):
        t, _ = common.count_targets_and_lures(seq, self.n)
        return 1.0 - (self.targets_norm.pdf(t) / self.targets_norm.pdf(self.targets))


if __name__ == '__main__':
    import time

    n = 3
    gen = SequenceGenerator(['1','2','3','4','5','6'], n)

    st = time.time()
    s = gen.generate(24, 4)
    st = time.time() - st
    print(f"{common.count_targets_and_lures(s,n)}")
    print(f"time = {st:0.2f}s")
    st = time.time()
    s = gen.generate(48, 16)
    st = time.time() - st
    print(f"{common.count_targets_and_lures(s,n)}")
    print(f"time = {st:0.2f}s")
    st = time.time()
    s = gen.generate(72, 8)
    st = time.time() - st
    print(f"{common.count_targets_and_lures(s,n)}")
    print(f"time = {st:0.2f}s")