from abc import abstractmethod
from dataclasses import dataclass
from itertools import product
from typing import Dict, Iterable, Set
from antlr4 import ParseTreeVisitor, InputStream, CommonTokenStream
from logic.grammars.OLexer import OLexer
from logic.grammars.OParser import OParser


class Expr:
    @abstractmethod
    def eval(self, assignment: Dict[str, bool]):
        pass


@dataclass
class AtomExpr(Expr):
    name: str

    def eval(self, assignment: Dict[str, bool]):
        return assignment[self.name]

@dataclass
class NotExpr(Expr):
    expr: Expr

    def eval(self, assignment: Dict[str, bool]):
        return not self.expr.eval(assignment)

@dataclass
class BinExpr(Expr):
    op: str
    left: Expr
    right: Expr

    def eval(self, assignment: Dict[str, bool]):
        if self.op == "|":
            return self.left.eval(assignment) or self.right.eval(assignment)
        if self.op == "&":
            return self.left.eval(assignment) and self.right.eval(assignment)
        if self.op == "->":
            return (not self.left.eval(assignment)) or self.right.eval(assignment)
        if self.op == "<->":
            return self.left.eval(assignment) == self.right.eval(assignment)


@dataclass
class Formula:
    alphabet: Set[str]
    root: Expr

    def models(self) -> Iterable[Iterable[str]]:
        alphabet = tuple(self.alphabet)
        for bools in product((True, False), repeat=len(alphabet)):
            assignment =  {atom: v for atom, v in zip(alphabet, bools)}
            if self.root.eval(assignment):
                yield assignment


class FormulaBuilder(ParseTreeVisitor):
    alphabet: Set[str]

    def __init__(self) -> None:
        self.alphabet = set()

    def visitParenth(self, ctx: OParser.ParenthContext):
        return self.visit(ctx.expr())

    def visit_bin_op(self, ctx):
        return BinExpr(ctx.op.text, self.visit(ctx.expr(0)), self.visit(ctx.expr(1)))

    def visitLog_neg(self, ctx:OParser.Log_negContext):
        return NotExpr(self.visit(ctx.expr()))

    def visitLog_or(self, ctx: OParser.Log_orContext):
        return self.visit_bin_op(ctx)

    def visitLog_impl(self, ctx: OParser.Log_implContext):
        return self.visit_bin_op(ctx)

    def visitLog_iff(self, ctx: OParser.Log_iffContext):
        return self.visit_bin_op(ctx)

    def visitLog_and(self, ctx: OParser.Log_andContext):
        return self.visit_bin_op(ctx)

    def visitAtom(self, ctx: OParser.AtomContext):
        text = ctx.getText()
        self.alphabet.add(text)
        return AtomExpr(text)

    def build(self, tree):
        expr = self.visit(tree)
        return Formula(self.alphabet, expr)


def parse(text: str):
    input_stream = InputStream(text)
    lexer = OLexer(input_stream)
    stream = CommonTokenStream(lexer)
    parser = OParser(stream)
    tree = parser.expr()
    formula = FormulaBuilder().build(tree)
    return formula