from typing import override

from ..agent.things import Thing

#= Formatting
def _balanced(s:str) -> bool:
    """ checks whether a string has balanced parantheses
    
    :param s: string to check
    :returns: True if string has balanced parantheses, else False
    """
    counter:int = 0
    for c in s:
        if c == "(":
            counter += 1
        elif c == ")":
            if counter <= 0:
                return False
            counter -= 1
    return counter == 0


def format(formula: str) -> str:
    """ formats a formula by wrapping in parantheses if not already.
        will not put parenthesis around symbols or their negation
        
        :param formula: string to format
        :returns: formatted string
    """
    if not len(formula) or formula.isalpha() or (
        formula[0] == "(" and formula[-1] == ")" and _balanced(formula[1:-1])
    ) or (formula[0] == "¬" and formula[1:].isalpha()):
        return formula
    else:
        return f"({formula})"
    

#= Classes
class Sentence(Thing):
    """ base class of a sentence, empty 
    
    A sentence is a logical statement that can be evaluated to be true or false. It can be composed of propositions and logical connectives (and, or, not, implies, iff). The base sentence represents an empty sentence that is always true and contains no propositions.
    """    
    @override
    def __repr__(self) -> str:
        """ String representation of a sentence is just the formula """
        return self.formula
    
    def evaluate(self, model: dict) -> bool:
        """ Evaluates a model over the sentence, returns True if model  satisfies sentence, else False.
        
        For a base Sentence, this is always True, as it represents an empty sentence.
        
        :param model: model to evaluate over
        :returns: True if model satisfies sentence, else False        
        """
        return True

    @property
    def formula(self) -> str:
        """ Returns the string representation of the sentence """
        return ""


    @property
    def operand(self) -> "Sentence":
        """ Returns the operand of the sentence, which is None for the base sentence, as it has no operand. This is used for negation, where the operand of a Not is the sentence it negates.
        
        :returns: None, as the operand of a base sentence is just itself
        """
        raise AttributeError(f"{type(self).__name__} has no operand, it is empty")


    @property
    def symbols(self) -> set["Proposition"]:
        """ Returns the set of symbols in the sentence
        
        Base sentence contains no symbols, so returns empty set.
        
        :returns: set of symbols in sentence
        """
        return set()
    
    
    @override
    def __eq__(self, x:"Sentence") -> bool:
        """ Checks if two sentences are the same, returns True if they are the same sentence, else False
        
        :param x: sentence to compare to
        :returns: True if x is the same sentence as self, else False
        """
        return (self is x)
    
    
    @override
    def __hash__(self) -> int:
        """ Creates unique hash for a proposition based on its name, so that propositions with the same name are considered the same proposition in sets and dictionaries 
        
        :returns: hash of the proposition based on its name
        """
        return hash(self.formula)
    
    
    def __bool__(self) -> None:
        """ Defines behaviour of bool() on a Proposition. Prevents evaluation of boolean logic using, e.g. and and or, to avoid confusion.
        
        :raises TypeError: if bool() is called on a Proposition, as this is not a valid operation
        """
        raise TypeError(f"cannot evaluate {type(self)}: use Proposition.evaluate(model)")


    def __call__(self, model:dict["Sentence", bool]) -> bool:
        """ Defines call behaviour of a proposition, so that it can be called as a function to evaluate it over a model. This allows for more intuitive evaluation of propositions, e.g. P(model) instead of P.evaluate(model).
        
        :param model: model to evaluate over
        :returns: True if model satisfies proposition, else False
        """
        return self.evaluate(model)
    
                 
    def __contains__(self, x:"Sentence") -> bool:
        """ Defines "in" behaviour for a sentence, so that it can be used to check if a proposition is in another sentence. For a base sentence, this is just checking if x is self, as the base sentence contains no symbols and cannot contain any propositions.
        
        :param x: propositional symbol to check for
        :returns: True if x is self, else False
        :raises TypeError: if x is not a Proposition or Not, as this is not a valid input for checking if it is in self
        """
        if x == self:
            return True
        if isinstance(x, Proposition):
            return (x in self.symbols)
        elif isinstance(x, Not) and isinstance(x.operand, Proposition):
            return (x.operand in self.symbols)
        else:
            raise TypeError(
                f"Cannot evaluate if {type(x).__name__} is in {type(self).__name__}")
    
    
    def __invert__(self) -> "Sentence":
        """ Defines behaviour of ~ operator, returning negation of proposition (¬self)
        
        :returns: Not(self), the negation of self
        """
        return Not(self) if self.symbols else self
    
    
    def __and__(self, x:"Sentence") -> "Sentence":
        """ Defines behaviour of & operator, returning conjunction (self ∧ input) 
        
        :param x: proposition to conjoin with self
        :returns: Conjunction(self, x), the conjunction of self and x
        """
        return Conjunction(self, x) if self.symbols else x
    
    
    def __or__(self, x:"Sentence") -> "Sentence":
        """ Defines behaviour of | operator, returning disjunction (self ∨ input) 
        
        :param x: proposition to disjoin with self
        :returns: Disjunction(self, x), the disjunction of self and x
        """
        return Disjunction(self, x) if self.symbols else x

    
    def onlyif(self, x:"Sentence") -> "Sentence":
        """ 'self only if input', converse of 'input if self' returns implication (self -> input)
        
        :param x: proposition that self only if x
        :returns: Implication(self, x), the implication of self and x
        """
        return Implication(self, x) if self.symbols else x
    
    
    def implies(self, x:"Sentence") -> "Sentence":
        """ Alias for onlyif, 'self implies input', returns implication 
        (self -> input)
        
        :param x: proposition that self implies
        :returns: Implication(self, x), the implication of self and x
        """
        return Implication(self, x) if self.symbols else x


    def iff(self, x:"Sentence") -> "Sentence":
        """ 'if and only if' , returns biconditional (self <-> input)
        
        :param x: proposition that self is equivalent to
        :returns: Biconditional(self, x), the biconditional of self and x
        """
        return Biconditional(self, x) if self.symbols else x



class Proposition(Sentence):
    """ Logical proposition representing logical statement """
    def __init__(self, name:str) -> None:
        """ Initializes a proposition with a name, which is used as the formula for the proposition. The name should be a string that uniquely identifies the proposition.
        
        :param name: string name of the proposition, used as the formula for the proposition
        """
        self.__name:str = name
        self.__operand = self


    @property
    @override
    def operand(self) -> "Proposition":
        """ Returns the operand of the proposition, which is just itself for a base proposition. This is used for negation, where the operand of a Not is the proposition it negates.
        
        :returns: self, as the operand of a base proposition is just itself
        """
        return self.__operand
    

    @property
    @override
    def formula(self) -> str:
        """ the string representation of the proposition 
        
        :returns: the name of the proposition as its string representation
        """
        return self.__name
   
    
    @property
    @override
    def symbols(self) -> set["Proposition"]:
        """ set of logical symbols in proposition. For a base proposition, this is just itself as a symbol.
        
        :returns: set containing self, as a proposition is a symbol
        """
        return {self}
    
    
    @override
    def evaluate(self, model: dict[Sentence, bool]) -> bool:
        """ Evaluates a model containing self. Presumes that the model contains self as a key with a boolean value, and returns that value. If self is not in the model, raises an IndexError, as this is an invalid model for evaluating the proposition.
        
        :param model: model to evaluate over, should contain self as a key with a boolean value
        :returns: True if model satisfies proposition, else False
        :raises IndexError: if self is not in model, as this is an invalid model for evaluating the proposition
        """
        if self in model:
            return bool(model[self])
        else:
            raise IndexError(f"{self} not in model {model}")

                             
class Not(Sentence):
    """ Logical proposition representing negation or 'not'
    
    Created by applying the ~ operator to a proposition, e.g. ~P creates Not(P), which represents the negation of P (¬P). Not can only be applied to propositions, and represents the logical negation of the proposition it is applied to.
    """
    def __init__(self, symbol:Sentence):
        """ Constructor for Not, takes a proposition as input and creates the negation of that proposition.
        
        :param symbol: proposition to negate, should be a Proposition
        """
        # super().__init__("¬"+format(symbol.formula))
        self.__operand:Sentence = symbol
    
    
    @property
    @override
    def operand(self) -> Sentence:
        """ Returns the operand of the proposition, which is just itself for a base proposition. This is used for negation, where the operand of a Not is the proposition it negates.
        
        :returns: self, as the operand of a base proposition is just itself
        """
        return self.__operand
    
    
    @property
    @override
    def formula(self) -> str:
        """ String representation of the negation, which is just "¬" followed by the formatted formula of the operand.
        
        :returns: string representation of the negation
        """
        return "¬" + format(self.operand.formula)
    
    
    @override
    def evaluate(self, model:dict[Sentence, bool]) -> bool:
        """ Evaluates a model containing operand 
        
        :param model: model to evaluate over, should contain operand as a key with a boolean value
        :returns: True if model satisfies negation of operand, else False
        :raises IndexError: if operand is not in model, as this is an invalid model for evaluating the negation
        """
        return not self.operand.evaluate(model)


    @property
    @override
    def symbols(self) -> set[Proposition]:
        """ Returns the symbols in the negation, which are just the symbols in the operand, as the negation itself is not a symbol.
        
        :returns: set of symbols in the operand, as the negation itself is not a symbol
        """
        return self.operand.symbols


    @override
    def __contains__(self, x:Proposition) -> bool:
        """ Checks if a proposition is in the negation, which is just checking if it is in the operand, as the negation itself is not a symbol and cannot contain any propositions.
        
        :param x: proposition to check for
        :returns: True if x is in operand, else False, as the negation itself is not a symbol and cannot contain any propositions
        """
        if type(x) is Not and type(x.operand) is Proposition:
            return x.operand == self.operand
        return x in self.operand


    @override
    def __invert__(self) -> Sentence:
        """ returns ¬(¬operand), so operand"""
        return self.operand


class Conjunction(Sentence):
    """ Logical conjunction (∧) of two or more propositions """
    def __init__(self, operand0:Sentence, operand1:Sentence, *operands):
        """ Create a conjunction of two or more Propositions
        
        :param operand0: first proposition to conjoin
        :param operand1: second proposition to conjoin
        :param operands: additional propositions to conjoin, if any
        """
        _operands:list = self._expand(operand0) + self._expand(operand1)
        for operand in operands:
            _operands.extend(self._expand(operand))

        # formula:str = format(_operands[0].formula)
        # for operand in _operands[1:]:
        #     formula += " ∧ " + format(operand.formula)

        # super().__init__(formula)
        self.operands:tuple = tuple(_operands)


    @property
    @override
    def formula(self) -> str:
        """ Returns the string representation of the conjunction, which is just the formula of each operand joined by " ∧ " 
        
        :returns: string representation of the conjunction
        """
        formula:str = format(self.operands[0].formula)
        for operand in self.operands[1:]:
            formula += " ∧ " + format(operand.formula)
        return formula


    def _expand(self, operand:Sentence) -> list[Sentence]:
        """ returns individual operands in a conjunction
            (P ∧ Q) becomes [P, Q]
            R becomes [R]
        """
        if isinstance(operand, Conjunction):
            return list(operand.operands)
        else:
            return [operand]
    
    
    @override
    def __and__(self, operand: Sentence) -> "Conjunction":
        """ As self is conjunction, adding another operand with & should just add it to the list of operands, so returns a new conjunction with the new operand added to the list of operands.
        
        :param operand: proposition to add to conjunction
        :returns: new Conjunction with operand added to list of operands
        """
        return Conjunction(*self.operands, operand)


    @override
    def evaluate(self, model:dict[Sentence, bool]) -> bool:
        """ Evanluates a model containing all operands, returns True if model satisfies conjunction, else False
        
        :param model: model to evaluate over, should contain all operands as keys with boolean values
        :returns: True if model satisfies conjunction, else False
        """
        return all(operand.evaluate(model) for operand in self.operands)
    
    
    @override
    def __contains__(self, x:Sentence) -> bool:
        """ Checks if a proposition is in the conjunction, which is just checking if it is in any of the operands, as the conjunction itself is not a symbol and cannot contain any propositions. 
        
        :param x: proposition to check for
        :returns: True if x is in any operand, else False, as the conjunction itself
        """
        return any(x in operand for operand in self.operands)


    @property
    @override
    def symbols(self) -> set[Proposition]:
        """ Returns the symbols in the conjunction, which are just the symbols in all of the operands, as the conjunction itself is not a symbol.
        """
        return set.union(*[operand.symbols for operand in self.operands])

And = Conjunction  # alias


class Disjunction(Sentence):
    """ Logical disjunction (∨) of two or more propositions """
    def __init__(self, operand0:Sentence, operand1:Sentence, *operands):
        """ Create a disjunction of two or more Propositions
    
        :param operand0: first proposition to disjoin
        :param operand1: second proposition to disjoin
        :param operands: additional propositions to disjoin, if any
        """
        _operands:list = self._expand(operand0) + self._expand(operand1)
        for operand in operands:
            _operands.extend(self._expand(operand))

        # formula:str = format(_operands[0].formula)
        # for operand in _operands[1:]:
        #     formula += " ∨ " + format(operand.formula)

        # super().__init__(formula)
        self.operands:tuple = tuple(_operands)
    
    @property
    @override
    def formula(self) -> str:
        """ Returns the string representation of the disjunction, which is just the formula of each operand joined by " ∨ " 
        
        :returns: string representation of the disjunction
        """
        formula:str = format(self.operands[0].formula)
        for operand in self.operands[1:]:
            formula += " ∨ " + format(operand.formula)
        return formula
    
    
    def _expand(self, operand:Sentence) -> list[Sentence]:
        """ returns individual operands in a disjunction
            (P ∨ Q) becomes [P, Q]
            R becomes [R]
            
            :param operand: proposition to expand
            :returns: list of propositions that are the individual operands in the disjunction
        """
        if isinstance(operand, Disjunction):
            return list(operand.operands)
        else:
            return [operand]
    
    @override
    def __or__(self, operand: Sentence) -> "Disjunction":
        """ As self is disjunction, adding another operand with | should just add it to the list of operands, so returns a new disjunction with the new operand added to the list of operands.
        
        :param operand: proposition to add to the disjunction
        :returns: new disjunction with the new operand added
        """
        return Disjunction(*self.operands, operand)
    
    
    @override
    def evaluate(self, model:dict[Sentence, bool]) -> bool:
        """ Evaluates a model containing all operands, returns True if model satisfies disjunction, else False
        
        :param model: model to evaluate over, should contain all operands as keys with boolean values
        :returns: True if model satisfies disjunction, else False
        """
        return any(operand.evaluate(model) for operand in self.operands)
    
    
    @override 
    def __contains__(self, x:Sentence) -> bool:
        """ Checks if a proposition is in the disjunction, which is just checking if it is in any of the operands, as the disjunction itself is not a symbol and cannot contain any propositions.
        
        :param x: proposition to check for
        :returns: True if x is in any operand, else False, as the disjunction itself
        """
        return any(x in operand for operand in self.operands)


    @property
    @override
    def symbols(self) -> set[Proposition]:
        """ Return the symbols in the disjunction, which are just the symbols in all of the operands, as the disjunction itself is not a symbol.
        
        :returns: set of symbols in the disjunction
        """
        return set.union(*[operand.symbols for operand in self.operands])
    
    
Or = Disjunction  # alias


class Implication(Sentence):
    """ Logical implication (->) of two propositions """
    def __init__(self, left:Sentence, right:Sentence):
        """ Create an implication of two propositions, left and right, representing left -> right
        
        :param left: proposition representing the antecedent of the implication
        :param right: proposition representing the consequent of the implication
        """
        self.antecedent:Sentence = left
        self.consequent:Sentence = right
        self.left:Sentence = self.antecedent
        self.right:Sentence = self.consequent


    @property
    @override
    def formula(self) -> str:
        """ String representation of the implication, formatted as "antecedent -> consequent" with appropriate formatting for the antecedent and consequent.
        
        :returns: string representation of the implication
        """
        return f"{format(self.antecedent.formula)} -> {format(self.consequent.formula)}"
    
    
    @override
    def evaluate(self, model:dict[Sentence, bool]) -> bool:
        """ Evaluates a model containing the antecedent and consequent, returns True if model satisfies implication, else False
        
        :param model: model to evaluate over, should contain antecedent and consequent as keys with boolean values
        :returns: True if model satisfies implication, else False
        """
        return ((not self.antecedent.evaluate(model)) or 
                    self.consequent.evaluate(model))


    @property
    @override
    def symbols(self) -> set[Proposition]:
        """ Returns the symbols in the implication, which are just the symbols in the antecedent and consequent, as the implication itself is not a symbol.
        
        :returns: set of symbols in the implication, which are just the symbols in the antecedent and consequent
        """
        return set.union(self.antecedent.symbols, self.consequent.symbols)
    

class Biconditional(Implication):
    """ Logical biconditional (<->) of two propositions 
    
    Expands functionality of Implication, as a biconditional is just two implications (P <-> Q is equivalent to (P -> Q) ∧ (Q -> P)). Therefore, Biconditional inherits from Implication, and overrides the formula and evaluate methods to reflect the biconditional nature of the statement.
    """
    @property
    @override
    def formula(self) -> str:
        """ Formula representation of the biconditional, formatted as "left <-> right" with appropriate formatting for left and right.
        """
        return f"{format(self.left.formula)} <-> {format(self.right.formula)}"
    
    
    @override
    def evaluate(self, model: dict[Sentence, bool]) -> bool:
        """ Evaluates a model containing the left and right propositions, returns True if model satisfies biconditional, else False
        
        :param model: model to evaluate over, should contain left and right propositions as keys with boolean values
        :returns: True if model satisfies biconditional, else False
        """
        left = self.left.evaluate(model)
        right = self.right.evaluate(model)
        return (left and right) or (not left and not right)
    

#= Utility functions
def simplify(statement:Sentence) -> Sentence:
    """ Simplifies sentences into logical ands and ors
    
    :param statement: sentence to simplify, can be any sentence composed of the logical connectives and sentences defined above
    :returns: simplified sentence, where all implications and biconditionals have been rewritten as conjunctions and disjunctions, and all negations have been pushed down to the level of sentences
    """
    match statement:
        case Biconditional():
            l, r = statement.antecedent, statement.consequent
            return simplify(l.implies(r) & r.implies(l))
        case Implication():
            return simplify(~statement.antecedent | statement.consequent)
        case Not():
            proposition = simplify(statement.operand)
            match proposition:
                case Conjunction() | Disjunction():
                    return Conjunction(*[~p for p in proposition.operands])
                case _:
                    return ~proposition
        case Conjunction():
            out = Conjunction(
                *[simplify(p) for p in statement.operands if p is not None])
            return out
        case Disjunction():
            out = Disjunction(
                *[simplify(p) for p in statement.operands if p is not None])
            return out
        case _:
            return statement


def distribute(sentence:Sentence) -> Sentence:
    """ Distributes disjunctions 
    
    :param sentence: sentence to distribute, should be a disjunction of sentences and conjunctions
    :returns: sentence with disjunctions distributed, so that all conjunctions are at the lowest level of the sentence, and all disjunctions are at the highest level of the sentence, i.e. in conjunctive normal form
    """
    if type(sentence) is Conjunction:
        return Conjunction(
            *[distribute(operand) for operand in sentence.operands])
    if type(sentence) is not Disjunction:
        return sentence
    if not any([type(p) is Conjunction for p in sentence.operands]):
        return sentence
    distr = []
    for proposition in sentence.operands:
        if type(proposition) is Conjunction:
            for operand in proposition.operands:
                conj = Disjunction(operand,*[p for p in sentence.operands if p is not proposition])
                distr.append(conj)
    return Conjunction(*distr)


def cnf(sentence:Sentence) -> Sentence:
    """ returns the conjunctive normal form of a logical sentence
    
    :param sentence: sentence to convert to conjunctive normal form, should be a sentence composed of the logical connectives and sentences defined above
    :returns: sentence in conjunctive normal form, where all conjunctions are at the lowest level of the sentence, and all disjunctions are at the highest level of the sentence, i.e. in conjunctive normal form
    """
    return distribute(simplify(sentence))