import random
import scipy.stats

import benchmarks.common as common


class SequenceGenerator:
    """nb_gm_005:
        - pseudo-random sampling.
        - fixed number of targets,lures, and distractors.
        - all stimuli presented equally as target, lure, or distractor.

    """

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

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

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

        self.targets_ratio_norm = scipy.stats.norm(targets/trials, 0.2)
        self.lures_ratio_norm   = scipy.stats.norm(lures/trials, 0.2)
        self.skewness_norm      = scipy.stats.norm(0      , trials/len(self.choices))

        print(self.skewness_norm.pdf(0))
        print(self.targets_ratio_norm.pdf(1.0/3.0))
        print(self.lures_ratio_norm.pdf(1.0/6.0))

        self.seq = []

    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, n)
        targets_ratio, lures_ratio = targets/len(seq), lures/len(seq)
        c1, c2, c3 = self.skewness_cost(seq), self.targets_cost(targets_ratio), self.lures_cost(lures_ratio)
        # print(c1,c2,c3)
        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)

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

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


if __name__ == '__main__':
    import time
    import numpy as np

    trials, targets, lures = 24, 8, 4
    choices = ['1', '2', '3', '4', '5', '6', '7', '8']
    print('alg,n,trials,time,targets,lures')
    for s in range(1, common.sample_size):
        n = np.random.randint(2, 8)
        gen = SequenceGenerator(choices, n)
        st = time.time()
        s = gen.generate(trials, targets, lures)
        st = time.time() - st
        t, l = common.count_targets_and_lures(s, n)
        print(f"nb_gm_004,{n},{trials},{st:0.2f},{t},{l}")