import numpy as np
import random
from ...agent.things import Thing, Agent
from ...agent.environment import Environment
from collections.abc import Callable, Iterable
from copy import deepcopy

from typing import Literal, override, TypeVar, Collection, Generic, Iterator

from .util import aslist

__all__ = [
    'util',
    'CSPAgent',
    'ConstraintSatisfactionProblem',
    'CSPRunnerEnvironment',
    'Variable',
    'Factor']

# === CLASSES === #
Type = TypeVar("Type")
Domain = Collection[Type]

class __Variable(Thing, Generic[Type]):
    """ Abstract Base Class for Variables in a CSP.
    
    Defines operations on variables to allow arithmetic and comparison
    """
    @property
    def value(self) -> Type:
        """ The value assigned to the variable (read only) """
        return self.__value
    
    @property
    def is_assigned(self) -> bool:
        """ Whether the variable has been assigned a value. """
        return hasattr(self, '__value') and self.__value is not None

    @override
    def __eq__(self, x:Type) -> bool:
        """ Equality comparison operator override. """
        return (x == self.value)
    
    @override
    def __ne__(self, x:Type) -> bool:
        """ Inequality comparison operator override. """
        return (x != self.value)
    
    @override
    def __lt__(self, x:Type) -> bool:
        """ Less-than comparison operator override. 
        
        Defaults to True if variable is unassigned.
        """
        return self.value < x if self.is_assigned else True

    @override
    def __gt__(self, x:Type) -> bool:
        """ Greater-than comparison operator override.
        
        Defaults to True if variable is unassigned.
        """
        return self.value > x if self.is_assigned else True
    
    @override
    def __le__(self, x:Type) -> bool:
        """ Less-than-or-equal comparison operator override. """
        return self.__eq__(x) or self.__lt__(x)
    
    @override
    def __ge__(self, x:Type) -> bool:
        """ Greater-than-or-equal comparison operator override. """
        return self.__eq__(x) or self.__gt__(x)
    
    @override
    def __add__(self, x:Type) -> Type | "__Variable":
        """ Addition operator override. """
        return self.value + x if self.is_assigned else self
    
    @override
    def __sub__(self, x:Type) -> Type | "__Variable":
        """ Subtraction operator override. """
        return self.value - x if self.is_assigned else self

    @override
    def __mul__(self, x:Type) -> Type | "__Variable":
        """ Multiplication operator override. """
        return self.value * x if self.is_assigned else self
    
    @override
    def __truediv__(self, x:Type) -> Type | "__Variable":
        """ True division operator override. """
        return self.value / x if self.is_assigned else self
    
    @override
    def __rtruediv__(self, x:Type) -> Type | "__Variable":
        """ Right-hand true division operator override. """
        return x / self.value if self.is_assigned else self
    
    @override
    def __mod__(self, x:Type) -> Type | "__Variable":
        """ Modulo operator override. """
        return self.value % x if self.is_assigned else self

    @override
    def __or__(self, x:Type) -> Type | "__Variable":
        """ Or operator override. """
        return self.value | x if self.is_assigned else self
    
    @override
    def __and__(self, x:Type) -> Type | "__Variable":
        """ And operator override. """
        return self.value & x if self.is_assigned else self
    
    @override
    def __pow__(self, x:Type) -> Type | "__Variable":
        """ Power operator override. """
        return self.value ** x if self.is_assigned else self

    @override
    def __neg__(self) -> Type | "__Variable":
        """ Negation operator override. """
        return -self.value if self.is_assigned else self

    @override
    def __truediv__(self, x:Type) -> Type | "__Variable":
        """ True division operator override. """
        return self.value / x if self.is_assigned else self
    
    @override
    def __floordiv__(self, x:Type) -> Type | "__Variable":
        """ Floor division operator override. """
        return self.value // x if self.is_assigned else self
    
    @override
    def __abs__(self) -> Type | "__Variable":
        """ Absolute value operator override. """
        return abs(self.value) if self.is_assigned else self

class Variable(__Variable, Generic[Type]):
    """ Variable class for Constraint Satisfaction Problems. """
    def __init__(self, domain:Domain[Type], name:str | None = None):
        """ Initialise a variable with a domain and optional name.
        
        :param domain: Collection of possible values for the variable.
        :param name: Optional name for the variable.
        """
        self.__domain = deepcopy(domain)  # domain belongs only to one variable
        self.__value:Type | None = None
        self.name = name
        self.__hash = random.random()  # unique hash for variable identity

    @override
    def __hash__(self ) -> int:
        """ Returns a unique hash for the variable. """
        return hash(self.__hash)

    @property
    def is_assigned(self) -> bool:
        """ Determines whether variable is assigned a value. """
        return self.__value is not None

    @property
    def value(self) -> Type | None:
        """ Returns the value assigned to the variable  """
        return self.__value
    
    @value.setter
    def value(self, x: Type | None) -> None:
        """ Set value of variable to x if in domain, else raise ValueError.
        
        :param x: Value to assign to variable.
        :raises ValueError: If x is not in variable domain.
        """
        if x in self.domain or x is None:
            self.__value = x
        else:
            raise ValueError(f"{self.name}: {x} not in domain {self.domain}")

    @property
    def domain(self) -> Domain[Type]:
        """ Returns the domain of the variable. (read only)"""
        return self.__domain
    
    def __repr__(self) -> str:
        """ String representation of the variable. Returns name and either value if assigned or ?. """
        return f"{self.name}({str(self.value) if self.is_assigned else '?'})"


class Factor(Thing):
    """ Wrapper class for constraints in a CSP. """
    def __init__(self,
                 constraint:Callable[..., bool],
                 variables:Iterable[Variable] | Variable):
        """ Initialise a Factor with a constraint function and variables.
        
        :param constraint: A callable function representing the constraint.
        :param variables: An Variable (or iterable collection of Variable instances) involved in the constraint.
        """
        assert isinstance(constraint, Callable), "constraint must be callable"
        self.__function = constraint
        if not isinstance(variables, Iterable):
            variables = [variables]
        self.__variables:list[Variable] = aslist(variables)
        if self.__xnary < 1:
            raise ValueError(f"{self}: number of variables must be >1")

    def __call__(self, *args, **kwargs) -> bool:
        """ Implements ability to call that will evaluate the constraint function. """
        return self.__function(*args, **kwargs)

    def __iter__(self) -> Iterator[Variable]:
        """ Implements iteration over variables in the factor. """
        return iter(self.__variables)
    
    def __contains__(self, variable: Variable) -> bool:
        """ Overrides in keyword behaviour because variable equality 
         is checked by value rather than hash
         
        :param variable: Variable to check for membership in the factor.
        :return: True if variable is in the factor, else False.
        """
        return any(variable is _variable for _variable in self)
        
    @property
    def is_satisfied(self) -> bool:
        """ Determines whether the constraint is satisfied. """
        if all(v.is_assigned for v in self.__variables):
            return self(*self.__variables) 
        elif self.is_global:
            return self.__function(*self.__variables)
        else:
            return True
    
    @override
    def __repr__(self) -> str:
        """ String representation of the factor. """
        return str(tuple([str(v.name) for v in self.variables]))

    @property
    def __xnary(self) -> int:
        """ Determines xnary of the factor. """
        return len(self.variables)
    
    @property
    def is_unary(self) -> bool:
        """ Is the factor unary? """
        return self.__xnary == 1
    
    @property
    def is_binary(self) -> bool:
        """ Is the factor binary? """
        return self.__xnary == 2
    
    @property
    def is_global(self) -> bool:
        """ Is the factor global? """
        return self.__xnary > 2

    @property
    def arcs(self) -> tuple["Arc", "Arc"]:
        """ Returns the arcs for binary factors.
        
        :return: Tuple of arcs (Factor, Variable, Variable), e.g. (self, A, B) and (self, B, A)
        :raises TypeError: If factor is not binary.
        """
        if not self.is_binary:
            raise TypeError(f"{self}: constraint is not binary")
        L,R = self.variables
        return ((self, L, R), (self, R, L))

    @property
    def variables(self) -> Collection[Variable]:
        """ Returns the variables involved in the factor. (read only) """
        return self.__variables
    
Arc = tuple[Factor, Variable, Variable]  # type alias for arcs in binary factors

class Constraint(Factor):
    """ Class alias for Factor to represent constraints in a CSP. """
    pass

class ConstraintSatisfactionProblem:
    """ Constraint Satisfaction Problem base class.
    
    Defined by variables (and their domains) and constraints.
    
    Has attributes
    - variables: Collection of variables in the CSP.
    - domains: Dictionary of variable domains.
    - constraints: Collection of constraints (Factors) in the CSP.
    
    Has utility checks including
    - is_consistent: Whether all constraints are satisfied.
    - is_complete: Whether all variables are assigned and all constraints are satisfied.
    - arcs: List of arcs in binary constraints.
    """
    def __init__(self, 
                 variables:Collection[Variable],
                 constraints:Collection[Factor]) -> None:
        """ 
        Constructor for CSP
        
        :param variables: Collection of variables in the CSP.
        :param constraints: Collection of constraints (Factors) in the CSP.
        """
        self.__variables = variables
        self.__constraints = constraints

    @property
    def variables(self) -> Collection[Variable]:
        """ Read only list of variables in the CSP. 
        
        Read only.
        
        :return: List of variables.
        """
        return self.__variables
    
    @property
    def domains(self) -> dict[Variable, Domain]:
        """ Read only dictionary of variable domains. 
        
        :return: Dictionary mapping variables to their domains.
        """
        return {variable: variable.domain for variable in self.variables}

    @property
    def constraints(self) -> Collection[Factor]:
        """ Read only list of constraints in the CSP."""
        return self.__constraints
    
    @property
    def is_consistent(self) -> bool:
        """ Determines whether all constraints are satisfied for current variable assignments. """
        return all(constraint.is_satisfied for constraint in self.constraints)
    
    @property
    def is_complete(self) -> bool:
        """ Determines whether all variables are assigned and all constraints are satisfied. """
        return all(variable.is_assigned for variable in self.variables) \
                    and self.is_consistent
    
    @property
    def arcs(self) -> list[Arc]:
        """ Returns list of arcs in binary constraints. """
        return [arc for constraint in self.constraints if constraint.is_binary
                        for arc in constraint.arcs]
    
    def __repr__(self) -> str:
        """ String representation of the CSP. """
        return str({v.name: v for v in self.variables})
    

CSP = ConstraintSatisfactionProblem  # type alias

class CSPAgent(Agent):
    """ ConstraintSatisfaction Problem Agent base class. """
    def __repr__(self) -> str: 
        return "🤖"
        
    @override
    def program(self, percept:CSP) -> tuple[Literal["solve"], CSP]:
        """ Constraint Satisfaction Problem agent program. """
        return ("solve", percept)
    
    def solve(self, csp: CSP) -> CSP:
        """ Actuator for solving CSP problems """
        raise NotImplementedError("Need to implement solve method in subclass")


## Environment ##
class CSPRunnerEnvironment(Environment):
    """ Environment for running CSP agents. """
    def __init__(self, csp:CSP, solution=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.csp = csp
        self.solution = solution
        self.__done = False

    @property
    def is_done(self) -> bool:
        """ Check if CSP is solved. """
        return self.__done

    @override
    def percept(self, agent:CSPAgent) -> CSP:
        """ Provides agent with the CSP as percept. """
        return self.csp

    @override
    def execute_action(self,
                       agent:CSPAgent,
                       action:tuple[Literal["solve"], CSP]):
        """ Execute action for CSP agent. 
        
        :param agent: The CSP agent executing the action.
        :param action: The action to execute, expected to be a tuple ("solve", csp).
        """
        command, csp = action
        if command == "solve":
            solution = agent.solve(csp)
        else:
            raise ValueError(f"{agent}: unknown command {command}")
        
        if solution:
            print(f"final solution:\n{solution}")
            if self.solution is not None:
                print(f"correct solution:\n{self.solution}")
        else:
            print("no solution found")
        self.__done = True


    @override
    def run(self, steps:int=1, pause_for_user:bool=False) -> None:
        """ Run the CSP environment for a number of steps."""
        if pause_for_user:
            input("Press enter to start simulation")

        print(f"{self}: Running for {steps} iterations.")
        for i in range(steps):
            if self.is_done:
                if steps:
                    print(f"{self}: Simulation complete after {i} of {steps} iterations.")
                return
            self.step()
        if steps:
            print(f"{self}: Simulation complete after {steps} of {steps} iterations.")