from AST import ONegation, OFormula, OBinary, OAtom, OConst, KB, atom, Set
from more_itertools import powerset


def oformula_sat(o: OFormula, total_I: Set[atom]):
    match o:
        case OConst.TRUE:
            return True
        case OConst.FALSE:
            return False
        case OBinary("|", left, right):
            return oformula_sat(left, total_I) or oformula_sat(right, total_I)
        case OBinary("&", left, right):
            return oformula_sat(left, total_I) and oformula_sat(right, total_I)
        case OBinary("->", left, right):
            return (not oformula_sat(left, total_I)) or oformula_sat(right, total_I)
        case OAtom(name):
            return name in total_I
        case ONegation(of):
            return not oformula_sat(of, total_I)
        case _:
            raise Exception("Ontology error")


def objective_knowledge_consistent(kb: KB, T: Set[atom], F: Set[atom]) -> bool:
    "OB_{T, F} has at least one model"
    undefined_atoms = (kb.katoms | kb.oatoms) - (T | F)
    for I in powerset(undefined_atoms):
        total_I = T.union(I)
        if oformula_sat(kb.ont, total_I):
            return True
    return False


def objective_knowledge(
    kb: KB, T: Set[atom], F: Set[atom]
) -> tuple[bool, Set[atom], Set[atom]]:
    "OB_{T, F} |= a"
    undefined_atoms = (kb.katoms | kb.oatoms) - (T | F)
    entailed = kb.katoms
    entailed_false = kb.katoms
    consistent = False
    for I in powerset(undefined_atoms):
        total_I = T.union(I)
        if oformula_sat(kb.ont, total_I):
            consistent = True
            entailed = entailed.intersection(total_I)
            entailed_false = entailed_false.intersection(kb.katoms.difference(total_I))
    return consistent, T.union(entailed), F.union(entailed_false)


def objective_knowledge_brave(kb: KB, T: Set[atom], F: Set[atom]):
    "OB_{T, F} |/= -a"
    undefined_atoms = (kb.katoms | kb.oatoms) - (T | F)
    entailed = Set()
    entailed_false = Set()
    consistent = False
    for I in powerset(undefined_atoms):
        total_I = T.union(I)
        if oformula_sat(kb.ont, total_I):
            consistent = True
            entailed = entailed.union(total_I)
            entailed_false = entailed_false.union(kb.katoms.difference(total_I))
    return consistent, T.union(entailed), F.union(entailed_false)