from typing import override
from copy import copy

from ..agent.things import ModelBasedAgent

from .logic import (
    Sentence,
    Proposition,
    Not,
    Conjunction,
    Disjunction,
    cnf  
)


def is_literal(x: Sentence) -> bool:
    """ Checks if a sentence is a literal, i.e. P or ~P
    
    :param x: sentence to check
    :returns: True if sentence is a literal, else False
    """
    return type(x) in [Proposition, Not]


class KnowledgeBase(Sentence):
    """ Represents the knowledge base of an agent """
    def __init__(self, *args:Sentence) -> None:
        """ Initialises the knowledge base with a sentence or sequence of sentences. If no arguments are provided, initialises as an empty.
        
        :param args: a sequence of sentences to add to the knowledge base, or a single sentence, or no arguments for an empty knowledge base
        """
        match len(args):
            case 0:
                self.__sentence:Sentence = Sentence()
            case 1:
                if not isinstance(args[0], KnowledgeBase):
                    self.__sentence:Sentence = args[0]
                else:
                    self.__sentence:Sentence = args[0].sentence
            case _:
                if any(isinstance(arg, KnowledgeBase) for arg in args):
                    args = tuple(arg.sentence if isinstance(arg, KnowledgeBase) else arg for arg in args)
                self.__sentence:Sentence = Conjunction(*args)

        self._entails:set[Proposition|Not] = set()


    @override
    def __repr__(self) -> str:
        """ String representation of the knowledge base is just the string representation of the sentence it contains """
        return f"KB: {self.__sentence}"


    @property
    def sentence(self) -> Sentence:
        """ Returns the sentence contained in the knowledge base
        
        :returns: the sentence contained in the knowledge base
        """
        return self.__sentence


    @override
    def evaluate(self, model: dict[Sentence, bool]) -> bool:
        """ evalutes a model over the knowledge base 
        
        :param model: a dictionary mapping sentences to their truth values, representing a possible world or interpretation of the sentences in the knowledge base
        :returns: the truth value of the knowledge base in the given model, which is just the truth value of the sentence contained in the knowledge base evaluated in the model
        """
        return self.__sentence.evaluate(model)
    
    @override
    def __call__(self, model: dict[Sentence, bool]) -> bool:
        """ Alias for evaluate, allows the knowledge base to be called as a function with a model to evaluate it in
        """
        return self.evaluate(model)


    @property
    @override
    def symbols(self) -> set[Proposition]:
        """ Returns the set of symbols that exist in the knowledge base
        
        :return symbols: the set of symbols that exist in the knowledge base, which is just the set of symbols that exist in the sentence contained in the knowledge base
        """
        return self.__sentence.symbols


    def to_cnf(self) -> None:
        """ converts knowledge base into conjunctive normal form """
        self.__sentence = cnf(self.__sentence)
        
        
    def as_cnf(self) -> "KnowledgeBase":
        """ Returns a new knowledge base in conjunctive form 
        
        :returns: a new knowledge base in conjunctive normal form, which is just a copy of the current knowledge base with the sentence converted to conjunctive normal form
        """
        KB = copy(self)
        KB.to_cnf()
        return KB

    @override
    def __invert__(self) -> "KnowledgeBase":
        """ Apply negation over the knowledge base. Returns a new knowledge base representing the negation of the current knowledge base, which is just a copy of the current knowledge base with the sentence negated. If the knowledge base is empty, returns itself as the negation of an empty knowledge base is just an empty knowledge base.
        
        :param self: the knowledge base to negate
        :returns: a new knowledge base representing the negation of the current knowledge base, or itself if the knowledge base is empty
        """
        if len(self.symbols) > 0:
            return KnowledgeBase(~self.__sentence)
        else:
            return self
    
    
    @override
    def __and__(self, x:Sentence) -> Sentence:
        """ Defines & behaviour for knowledge base.
        
        :param x: the sentence to combine with the knowledge base
        :returns: a new sentence representing the conjunction of the knowledge base and the input sentence
        """
        if not isinstance(x, Sentence):
            raise TypeError(f"input must be Sentence, was type {type(x)}")
        if isinstance(self.__sentence, Sentence):
            return Conjunction(self.__sentence, x)
        else:
            return x


    @override
    def __or__(self, x:Sentence) -> Sentence:
        """ Defines | behaviour for knowledge base.
        
        :param x: the sentence to combine with the knowledge base
        :returns: a new sentence representing the disjunction of the knowledge base and the input sentence
        """
        if not isinstance(x, Sentence):
            raise TypeError(f"input must be Sentence, was type {type(x)}")
        if isinstance(self.__sentence, Sentence):
            return Disjunction(self.__sentence, x)
        else:
            return x


    @override
    def __contains__(self, x:Sentence) -> bool:
        """ Defines in behaviour by checking input sentence against KB sentence
        """
        return (x in self.__sentence)
    
    
    def __infer(self, x:Sentence) -> bool:
        """ applies inference algorithm to check if sentence is entailed by knowledge base.
        
        default behaviour uses model checking
        
        :param x: the sentence to check entailment of
        :returns: True if the knowledge base entails the input sentence, else False
        """
        return model_check(self, x)
    
    
    def entails(self, x:Sentence) -> bool:
        """ returns if knowledge base entails sentence
        
        :param x: the sentence to check entailment of
        :returns: True if the knowledge base entails the input sentence, else False 
        """
        if is_literal(x) and x in self._entails: return True
        
        entailment = self.__infer(x)
        
        if entailment and (type(x) in [Proposition, Not]):
            self._entails.add(x)  # add to set of literals for quick lookup
             
        return entailment
    
    
    def models(self, x:Sentence) -> bool:
        """ returns if knowledge base entails sentence (alias for self.entails)
        
        :param x: the sentence to check entailment of
        :returns: True if the knowledge base entails the input sentence, else False
        """
        return self.entails(x)


    def __add(self, x:Sentence, validate:bool = False) -> None:
        """ add a new proposition to the knowledge base
        
        :param x: the sentence to add to the knowledge base, should be a proposition or a conjunction of propositions
        :param validate: if True, checks if the new proposition contradicts the current knowledge base
        :raises ValueError: if the new proposition contradicts the current knowledge base (only if validate is True)
        """
        if not isinstance(x, Sentence):
            raise TypeError(f"input must be Sentence, was type {type(x)}")
        
        if len(self.symbols) > 0:  # non-empty knowledge base
            if validate and self.entails(~x):
                raise ValueError(
                    f"cannot add contradiction to Knowledge Base, KB ⊨ {~x}")
                return
            self.__sentence &= x  # set sentence conjunction of sentence and x
        else:
            self.__sentence = x  # add symbol for first time


    def add(self, *xs:Sentence, validate:bool = False) -> None:
        """ add new proposition(s) to the knowledge base 
        
        :param xs: one or more sentences to add to the knowledge base, or a single sentence, or no arguments for an empty knowledge base
        :param validate: if True, checks if the new propositions contradict the current knowledge base
        :raises ValueError: if any of the new propositions contradict the current knowledge base (only if validate is True)
        """
        for x in xs:
            self.__add(x, validate)  # add other arguments


## Inference function
def model_check(knowledge:KnowledgeBase, sentence:Sentence) -> bool:
    """ run model checking algorithm 
    
    :param knowledge: the knowledge base to check entailment against
    :param sentence: the sentence to check entailment of, should be a proposition or a conjunction of propositions
    :returns: True if the knowledge base entails the input sentence, else False
    """
    symbols = set.union(knowledge.symbols, sentence.symbols)
    return __check_all(knowledge, sentence, symbols, {})


def __check_all(knowledge:KnowledgeBase,
                query:Sentence,
                symbols:set[Proposition],
                model:dict[Sentence, bool]) -> bool:
    """ Utility function for model checking, recursively enumerates all models in knowledge base and  checks if query is entailed
    
    :param knowledge: the knowledge base to check entailment against
    :param query: the sentence to check entailment of
    :param symbols: the set of symbols to consider in the model
    :param model: the current model being evaluated
    :returns: True if the query is entailed by the knowledge base, else False
    """
    if not symbols:
        if knowledge.evaluate(model):
            return query.evaluate(model)
        return True
    else:
        remaining = symbols.copy()
        p = remaining.pop()  # random element from set

        model_true = model.copy()
        model_true[p] = True
        model_false = model.copy()
        model_false[p] = False
    return (__check_all(knowledge, query, remaining, model_true)
                and __check_all(knowledge, query, remaining, model_false))
    


class LogicalAgent(ModelBasedAgent):
    """ a logical agent with a knowledge base """
    def __init__(self, *args, **kwargs) -> None:
        """ Constructs a rational agent with an empty knowledge base. """
        super().__init__(*args, **kwargs)
        self.__knowledge:KnowledgeBase = KnowledgeBase()

    @property
    def knowledge(self) -> KnowledgeBase:
        """ Returns the knowledge base of the agent
        
        :returns: the knowledge base of the agent
        """
        return self.__knowledge