from core.explanations.utils import (
    model_to_assignment,
    list_to_bitmask
)
from pysat.solvers import Cadical195
from queue import Empty, Full
from tqdm import tqdm


class ExplanationGeneratorProducer:
    """Process to generate explanations using a SAT solver."""
    def __init__(self, batch_size, sample_size):
        self.batch_size = batch_size
        self.sample_size = sample_size

        # Initialize the SAT solver
        features = list(range(1, self.sample_size + 1))

        self.solver = Cadical195()
        self.solver.add_clause(features)

    def get_explanations(self, num_explanations, assignments, explanations_valid, explanations_invalid, use_progress_bar):
        explanations = set()

        # Number of assignments currently being processed
        processing = 0

        if use_progress_bar:
            progress_bar = tqdm(total=num_explanations, desc='Explanations', unit='explanation', leave=True)

        while len(explanations) < num_explanations:
            # Generate new assignments in batch
            for _ in range(self.batch_size - assignments.qsize()):
                # Check if there are assignments to explore
                satisfiable = self.solver.solve()

                if not satisfiable:
                    break

                # Get the model (assignment) from the solver
                model = self.solver.get_model()

                # Block the model to prevent re-exploration
                self.__block_model(model)

                try:
                    # Add the assignment to the queue
                    assignment = model_to_assignment(model)
                    assignments.put(assignment, timeout=0.5)
                    processing += 1
                except Full:
                    break

            # Block valid explanations
            while True:
                try:
                    explanation = explanations_valid.get(block=False)
                    self.__block_assignment_up(explanation)

                    explanation_bitmask = list_to_bitmask(explanation, self.sample_size)

                    if explanation_bitmask not in explanations:
                        explanations.add(explanation_bitmask)

                        if use_progress_bar:
                            progress_bar.update(1)

                    if len(explanation) == 0:
                        # If the explanation is empty, we can stop
                        satisfiable = False

                    processing -= 1
                except Empty:
                    break

            # Block invalid explanations
            while True:
                try:
                    explanation = explanations_invalid.get(block=False)
                    self.__block_assignment_down(explanation)

                    if len(explanation) == self.sample_size:
                        # If the explanation is the full assignment, we can stop
                        satisfiable = False

                    processing -= 1
                except Empty:
                    break

            if not satisfiable and processing <= 0:
                break

        if use_progress_bar:
            progress_bar.close()

        return explanations

    def __block_assignment_up(self, assignment):
        """Block an assignment and its supersets."""
        clause = [-(v+1) for v in assignment]
        self.solver.add_clause(clause)

    def __block_assignment_down(self, assignment):
        """Block an assignment and its subsets."""
        clause = [v+1 for v in range(self.sample_size) if v not in assignment]
        self.solver.add_clause(clause)

    def __block_model(self, model):
        """Block a model."""
        clause = [-v for v in model]
        self.solver.add_clause(clause)
