from core.explanations.generator.ExplanationGeneratorConsumer import ExplanationGeneratorConsumer
from core.explanations.generator.ExplanationGeneratorProducer import ExplanationGeneratorProducer
from multiprocessing import Event, Process, Queue


class ExplanationGenerator:
    def __init__(self, validator, ordering, num_consumers, sample_size):
        self.validator = validator
        self.ordering = ordering
        self.num_consumers = num_consumers
        self.sample_size = sample_size

    def get_explanations(self, sample, num_explanations, use_progress_bar):
        # Avoid repetitive work by setting the batch size to the number of consumers
        BATCH_SIZE = self.num_consumers

        # Use queues to communicate with the consumer processes
        assignments = Queue(maxsize=BATCH_SIZE)
        explanations_valid = Queue()
        explanations_invalid = Queue()

        # Start the consumer processes
        consumers = []
        stop_event = Event()

        for _ in range(self.num_consumers):
            features = self.validator.model.traversal_order(sample, self.ordering)

            consumer = ExplanationGeneratorConsumer(self.validator, features, stop_event)
            process = Process(target=consumer.explain, args=(sample, assignments, explanations_valid, explanations_invalid))
            process.start()
            consumers.append(process)

        producer = ExplanationGeneratorProducer(
            batch_size=BATCH_SIZE,
            sample_size=self.sample_size
        )

        explanations = producer.get_explanations(
            num_explanations=num_explanations,
            assignments=assignments,
            explanations_valid=explanations_valid,
            explanations_invalid=explanations_invalid,
            use_progress_bar=use_progress_bar
        )

        # Signal consumers to stop
        stop_event.set()

        # Wait for all consumer processes to stop
        for consumer in consumers:
            consumer.join()

        return explanations
