from constraint import *

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):
        seq = self._generate_initial_sequence()
        return self._optimize_sequence(seq)

    def _generate_initial_sequence(self):
        """
            Generates initial sequence of items based on choices and number of desired trials.
            In EvenRandom sequences, all stimuli have same number of appearances.
        """
        num_of_each_choice = int(self.trials / len(self.choices))
        seq = []
        for c in self.choices:
            seq.append([c for t in range(num_of_each_choice)])

        # if it's too short, add some new items
        if len(seq) < self.trials:
            seq += [self.choices[-1] for t in range(self.trials-len(seq))]

        return seq

    def _optimize_sequence(self, seq, tl_ratio):
        """Optimize a sequence to match a desired tl ratio"""

        p = Problem()

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

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

        return seq

    @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'])
    generated_seq = generator.generate()
    print('Even Random Sequence: %s' % ''.join(generated_seq))