from constraint import *

import itertools as it

import Numberjack as nj

class EvenRandomGenerator:
    """Generate even random sequences according to a predefined TL ration (Ralph, 2014)"""

    def __init__(self, choices, trials=64, tl=1):
        self.tl, self.trials, self.choices = tl, trials, choices

    def generate(self):
        seqs = self._generate_initial_sequences()
        return self._find_optimal_sequence(seqs, 0.2)

    def _generate_initial_sequences(self):
        """
            Generates initial sequence of items based on choices and number of desired trials.
            In EvenRandom sequences, all stimuli have same number of appearances.
        """

        pool = it.product(self.choices, repeat=self.trials)
        return pool

    def _find_optimal_sequence_numberjack(self,tl_ratio):
        """Optimize with Numberjack"""

        cost = nj.Variable(0, 100, 'cost')
        seqs = nj.Variable([]) # all sequences

        model = nj.Model(
            seqs.,
            cost == self.calculate_tl_ratio(seq) - tl_ratio,
            nj.Minimise(cost)
        )
        solver = model.load('Mistral')
        if solver.solve():
            solver.printStatistics()
        else:
            print("No solution with Numberjack")


    def _find_optimal_sequence(self, sequences, tl_ratio):
        """Optimize a sequence to match a desired tl ratio with python-constraints"""

        p = Problem()

        # TODO add all possible values for seq (its domain)
        p.addVariable("seq", list(sequences))

        p.addConstraint(lambda s: self.calculate_tl_ratio(s) - tl_ratio < 0.05)

        return p.getSolution()

    @staticmethod
    def calculate_tl_ratio(seq):
        """Calculates the T:L ratio of a sequence."""
        targets = 0
        lures = 0
        for index, item in seq:
            if item == seq[index-2]:
                targets += 1
            elif item == seq[index-1] or item == seq[index-3]:
                lures += 1
        # avoid division by zero
        if lures == 0:
            lures = 1
        return targets/lures


if __name__ == '__main__':
    generator = EvenRandomGenerator(['a', 'b', 'c'], trials = 4)
    generated_seq = generator.generate()
    print('Even Random Sequence: %s' % ''.join(generated_seq))