from collections.abc import Iterable
from collections import deque
from copy import deepcopy

from typing import TypeVar

# from ..csp import Variable, CSP, Factor # avoid circular imports
Variable = TypeVar("Variable")
CSP = TypeVar("CSP")
Factor = TypeVar("Factor")

import numpy as np

def alldiff(*variables:tuple[Variable|Iterable[Variable]]) -> bool:
    """ All different constraint function"""
    if len(variables) == 1:  # unwrap if given as a single iterable
        if not isinstance(variables[0], Iterable):
            return True
        variables = variables[0]
    # variables = [
    #     variable.value for variable in variable if variable.is_assigned]
    values = [
        variable.value for variable in variables 
            if variable.is_assigned]
    return len(set(values)) == len(values)  # all different


def aslist(npcol:Iterable) -> list:
    """ Enforces list type """
    return np.array(npcol).ravel().tolist()


def revise(factor:Factor, A:Variable, B:Variable) -> bool:
    """ Revise method for AC-3 algorithm.
    
    Checks every value in the domain of variable A to see if there is
    a corresponding value in the domain of variable B that satisfies. Removes any that do not.
    
    :param factor: The binary factor between variables A and B
    :param A: The first variable in the binary factor
    :param B: The other variable in the binary factor
    """
    if A.is_assigned: return False  # nothing to revise
    
    is_revised = False
    
    domain = deepcopy(A.domain)  # copy to avoid set size change errors
    for value in domain:  
        A.value = value  # temporarily assign A
        is_valid_B = False
        for _value in B.domain:
            B.value = _value  # temporarily assign B
            if factor.is_satisfied: is_valid_B = True
            B.value = None  # reset B
        if not is_valid_B:  # no valid B found for A=value
            A.domain.remove(value)  # remove value from A's domain
            is_revised = True
        A.value = None  # reset A
    return is_revised


def ac3(csp:CSP, log:bool=False, inplace:bool=True) -> CSP | bool | None:
    """ AC-3 algorithm for enforcing arc consistency on a CSP.
    
    :param csp: The constraint satisfaction problem to enforce arc consistency on
    :param log: If True, prints log messages during processing
    :param inplace: If True, modifies the given CSP, otherwise returns a new CSP
    :return: The modified CSP if inplace is False, True if inplace is True and successful
    """
    if not inplace: csp = deepcopy(csp)  # work on a copy if not inplace
    
    arcs = deque(csp.arcs)  # queue of arcs to consider
    
    while len(arcs) > 0:  # while there are arcs to consider
        f, A, B = arcs.popleft()  # end of queue
        if log:
            print(f"considering arc from {A.name} to {B.name}")
            print(f"  before: {A.name} in {A.domain}, {B.name} in {B.domain}")
        if revise(f, A, B):  # if we revised A's domain
            if log:
                print(
                    f"  after: {A.name} in {A.domain}, {B.name} in {B.domain}")
            if len(A.domain) == 0:  # domain wiped out, failure
                return False if inplace else None  # no possible solution
            
            for constraint in csp.constraints:  # update related arcs
                if log:
                    print(f"  do we need to check constraint {constraint}")
                    print(f"   is binary? {constraint.is_binary}")
                    print(f"   contains {A.name}? {A in constraint}")
                    print(f"   doesn't contain {B.name}? {B not in constraint}")
                if constraint.is_binary \
                        and A in constraint\
                            and B not in constraint:
                    for arc in constraint.arcs:
                        if arc[-1] is A:  # arc points to A
                            if arc not in arcs:  # not already in queue
                                if log:
                                    print(f"    adding {arc} to arcs frontier")
                                arcs.append(arc)  # add new arc to consider
                            elif log:
                                print(f"    arc {arc} already in frontier")
        elif log:
            print(f"  after: {A.name} in {A.domain}, {B.name} in {B.domain} (no change)")
    return True if inplace else csp


def make_node_consistent(csp, inplace=True) -> CSP | None:
    """ Makes the CSP node consistent by enforcing unary constraints.
    
    :param csp: The constraint satisfaction problem to make node consistent
    :param inplace: If True, modifies the given CSP, otherwise returns a new CSP
    :return: The modified CSP if inplace is False, otherwise None
    """
    if not inplace: csp = deepcopy(csp)
    
    for variable in csp.variables:
        if variable.is_assigned: continue  # ignore any assigned variables
        domain = variable.domain.copy()  # copy this to avoid set size change errors 
        for value in domain:
            variable.value = value
            for constraint in csp.constraints:
                if constraint.is_unary and variable in constraint:
                    if not constraint.is_satisfied:
                        variable.domain.remove(value)
                        break
            variable.value = None
    if not inplace: return csp