from queue import Queue
from typing import Any


class CSP:
    def __init__(
        self,
        variables: list[str],
        domains: dict[str, set],
        edges: list[tuple[str, str]],
    ):
        """Constructs a CSP instance with the given variables, domains and edges.

        Parameters
        ----------
        variables : list[str]
            The variables for the CSP
        domains : dict[str, set]
            The domains of the variables
        edges : list[tuple[str, str]]
            Pairs of variables that must not be assigned the same value
        """
        self.variables = variables
        self.domains = domains

        # Binary constraints as a dictionary mapping variable pairs to a set of value pairs.
        #
        # To check if variable1=value1, variable2=value2 is in violation of a binary constraint:
        # if (
        #     (variable1, variable2) in self.binary_constraints and
        #     (value1, value2) not in self.binary_constraints[(variable1, variable2)]
        # ) or (
        #     (variable2, variable1) in self.binary_constraints and
        #     (value1, value2) not in self.binary_constraints[(variable2, variable1)]
        # ):
        #     Violates a binary constraint
        self.binary_constraints: dict[tuple[str, str], set] = {}
        for variable1, variable2 in edges:
            self.binary_constraints[(variable1, variable2)] = set()
            for value1 in self.domains[variable1]:
                for value2 in self.domains[variable2]:
                    if value1 != value2:
                        self.binary_constraints[(variable1, variable2)].add(
                            (value1, value2)
                        )
                        self.binary_constraints[(variable1, variable2)].add(
                            (value2, value1)
                        )

    def ac_3(self) -> bool:
        """Performs AC-3 on the CSP.
        Meant to be run prior to calling backtracking_search() to reduce the search for some problems.

        Returns
        -------
        bool
            False if a domain becomes empty, otherwise True
        """

        def revise(csp, xi, xj) -> bool:
            """Reduce the domain of xi if possible

            Returns
            -------
            bool
                Whether or not there was a domain reduction
            """
            revised = False
            for x in set(csp.domains[xi]):
                if not any(  # translation of algorithm's forall expression.
                    [
                        (x, y) in csp.binary_constraints[(xi, xj)]
                        for y in csp.domains[xj]
                    ]
                ):
                    csp.domains[xi].remove(x)
                    revised = True
            return revised

        queue = Queue()
        for edge in self.binary_constraints.keys():
            queue.put(edge)

        while not queue.empty():
            (xi, xj) = queue.get()

            if revise(self, xi, xj):
                if len(self.domains[xi]) == 0:
                    return False
                for neighboring_edge in [
                    (a, b)
                    for (a, b) in self.binary_constraints.keys()
                    if a != xj and b == xi
                ]:
                    queue.put(neighboring_edge)
        return True

    def backtracking_search(self) -> None | dict[str, Any]:
        """Performs backtracking search on the CSP.

        Returns
        -------
        None | dict[str, Any]
            A solution if any exists, otherwise None
        """

        def backtrack(csp, assignment: dict[str, Any]) -> dict[str, Any] | None:
            """The recursive engine of backtracking_search.

            Returns
            -------
            dict[str, Any]
                The solution to bubble up
            None
                If current branch has no solution
            """
            # this is read from stdout in sudoku.py. important to 2c).
            # yes, it couldv'e been done differently, but this is
            # quick and unintrusive.
            print("i have been called")
            if len(assignment) == len(csp.variables):  # completion condition
                return assignment  # base-case
            var = select_unassigned_variable(csp, assignment)
            for value in order_domain_values(csp, var, assignment):
                if not consistent(csp, var, value, assignment):
                    continue  # do nothing (skip iteration)
                assignment[var] = value
                if result := backtrack(csp, assignment):
                    return result  # return if not failure
                assignment.pop(var)
            print("i have failed")
            return None  # failure

        def consistent(csp, var, value, assignment) -> bool:
            """Helper function to determine if an arc is consistent"""
            for v in assignment:
                if (var, v) in csp.binary_constraints.keys():
                    if (value, assignment[v]) not in csp.binary_constraints[(var, v)]:
                        return False
                if (v, var) in csp.binary_constraints.keys():  # bi-birectional
                    if (value, assignment[v]) not in csp.binary_constraints[(v, var)]:
                        return False
            return True

        def select_unassigned_variable(csp, assignment) -> str:
            """Helper function to find a variable that is yet to be assigned."""
            for v in csp.variables:
                if v not in assignment.keys():
                    return v
            return "this shouldn't happen"

        def order_domain_values(csp, var, assignment) -> list[Any]:
            """Helper function to order the domain values.
            Orders on the least constrained value
            """

            # calculate a sorting key by counting constraints
            def compare(value) -> int:
                s = 0
                for neighbor in [
                    b for (a, b) in csp.binary_constraints.keys() if a == var
                ]:
                    if neighbor in assignment.keys():
                        continue
                    for neighbor_value in csp.domains[neighbor]:
                        s += (value, neighbor_value) not in csp.binary_constraints[
                            (var, neighbor)
                        ]
                return s

            return sorted(list(csp.domains[var]), key=compare)

        return backtrack(self, {})  # start with no assignments


def alldiff(variables: list[str]) -> list[tuple[str, str]]:
    """Returns a list of edges interconnecting all of the input variables

    Parameters
    ----------
    variables : list[str]
        The variables that all must be different

    Returns
    -------
    list[tuple[str, str]]
        List of edges in the form (a, b)
    """
    return [
        (variables[i], variables[j])
        for i in range(len(variables) - 1)
        for j in range(i + 1, len(variables))
    ]