import random
import csv

import heapq


class SequenceGenerator:
    """Generate N-Back sequences with random sampling, but increase matching probability over time."""

    def __init__(
        self,
        choices: list,
        n=2,
        trials=24,
        target_probability_start=0.33,
        target_probability_end=0.50
    ):
        self.n, self.choices, self.trials = n, choices, trials
        self.target_probability_step = target_probability_end - target_probability_start
        self.target_probability = target_probability_start
        self.seq = []

    def generate(self) -> list:
        self.seq = []
        for t in range(self.trials):
            self.seq += self.random_sample()
        return self.seq

    def random_sample(self):
        is_target = (random.random() > self.target_probability)
        self.target_probability += self.target_probability_step
        choices = [item for item in self.choices if len(self.seq)<self.n or item != self.seq[-self.n]]
        return random.choice(choices) if len(self.seq) < self.n or not is_target else self.seq[-self.n]

    def count_targets_and_lures(self):
        n = self.n
        seq = self.seq
        targets = 0
        lures = 0
        for index in range(n, len(seq)):
            if seq[index] == seq[index - n]:
                targets += 1
            elif seq[index] == seq[index - (n-1)] or seq[index] == seq[index - (n+1)]:
                lures += 1
        return targets, lures