import fol


class Counter:
    def __init__(self):
        self.values = {}
        self.counter = 0

    def get(self, name):
        if name in self.values.keys():
            return self.values[name]
        self.values[name] = self.counter
        self.counter += 1
        return self.counter - 1

    def begin(self):
        self.values = {}


def elim_imp(f):
    res = fol.node.Node(f.name, f.value)

    if f.name == 'IMP':
        a = fol.node.Node('NOT')
        a.add_child(elim_imp(f.children[0]))
        res = fol.node.Node('OR')
        res.add_child(elim_imp(a))
        res.add_child(elim_imp(f.children[1]))
        return res
    else:
        for c in f.children:
            res.add_child(elim_imp(c))

        return res


def neg(f):
    res = f.copy()

    if f.name == 'NOT' and f.children[0].name in ['OR', 'AND']:
        p = fol.node.Node('NOT')
        p.add_child(neg(f.children[0].children[0]))
        q = fol.node.Node('NOT')
        q.add_child(neg(f.children[0].children[1]))

        if f.children[0].name == 'OR':
            res = fol.node.Node('AND')
            res.add_child(neg(p))
            res.add_child(neg(q))
            return res
        else:
            res = fol.node.Node('OR')
            res.add_child(neg(p))
            res.add_child(neg(q))
            return res

    res = fol.node.Node(f.name, f.value)

    for c in f.children:
        res.add_child(neg(c))

    return res


def normalize(f, counter):
    res = fol.node.Node(f.name, f.value)

    if f.name == 'VAR':
        res.value = f'x{counter.get(f.value)}'
        return res

    for child in f.children:
        res.add_child(normalize(child, counter))

    return res


def prenex(f):
    def mkneg(n):
        res = fol.node.Node('NOT')
        res.add_child(n)
        return res

    def inv(n):
        if n.name == 'FORALL':
            res = n.copy()
            res.name = 'EXISTS'
            return res
        elif n.name == 'EXISTS':
            res = n.copy()
            res.name = 'FORALL'
            return res
        else:
            return n.copy()

    res = fol.node.Node(f.name, f.value)

    if f.name == 'NOT':
        arg = f.children[0]

        if arg.name in ['FORALL', 'EXISTS']:
            res = inv(arg)
            res.children[1] = mkneg(arg.children[1])
            return res

    elif f.name in ['AND', 'OR', 'IMP']:
        def apply_prenex(lhs, rhs, op_name):
            swapped = False

            if rhs.name in ['FORALL', 'EXISTS']:
                lhs, rhs = rhs, lhs
                swapped = True

            if lhs.name in ['FORALL', 'EXISTS']:
                left = lhs.children[1]
                right = rhs
                if swapped:
                    left, right = right, left
                var = lhs.children[0]

                op = fol.node.Node(op_name)
                op.add_child(left)
                op.add_child(right)

                res = fol.node.Node(lhs.name)
                if op_name == 'IMP' and not swapped:
                    res.name = 'EXISTS' if lhs.name == 'FORALL' else 'FORALL'
                res.add_child(var)
                res.add_child(op)
                return res
            return f.copy()

        lhs = prenex(f.children[0])
        rhs = prenex(f.children[1])

        return apply_prenex(lhs, rhs, f.name)
    return f.copy()


def skolem(f):
    def collect_exists(n):
        res = []
        if n.name == 'EXISTS':
            res.append(n)
        for child in n.children:
            res.extend(collect_exists(child))
        return res

    def context(n, exists_quant, ctx=[]):
        if n.equals(exists_quant):
            return ctx

        if n.name == 'FORALL':
            ctx.append(n.children[0].value)

        for child in n.children:
            context(child, target, ctx)

        return ctx

    res = f.copy()
    targets = collect_exists(res)

    for target in targets:
        idx = targets.index(target)
        var_name = target.children[0].value

        subst = {
            var_name: fol.term(f'S{idx}')
        }

        ctx = context(f, target)

        if len(ctx) > 0:
            args = ','.join(ctx)
            subst = {
                var_name: fol.term(f's{idx}({args})')
            }

        res = res.subst(subst)

    return res.remove_quant('EXISTS')


def distrib(f):
    if len(f.children) != 2:
        return f

    lhs = f.children[0]
    rhs = f.children[1]

    if f.name == 'OR' and rhs.name == 'AND':
        a = distrib(lhs)
        b = distrib(rhs.children[0])
        c = distrib(rhs.children[1])
        res = fol.node.Node('AND')
        left = fol.node.Node('OR')
        left.add_child(a)
        left.add_child(b)
        right = fol.node.Node('OR')
        right.add_child(a)
        right.add_child(c)
        res.add_child(distrib(left))
        res.add_child(distrib(right))
        return res

    if f.name == 'OR' and lhs.name == 'AND':
        a = distrib(rhs)
        b = distrib(lhs.children[0])
        c = distrib(lhs.children[1])
        res = fol.node.Node('AND')
        left = fol.node.Node('OR')
        left.add_child(b)
        left.add_child(a)
        right = fol.node.Node('OR')
        right.add_child(c)
        right.add_child(a)
        res.add_child(distrib(left))
        res.add_child(distrib(right))
        return res

    return f


def cnf(f, counter=None):
    if counter is None:
        counter = Counter()

    g = distrib(
        skolem(prenex(normalize(neg(elim_imp(f)), counter)))
    ).remove_quant('FORALL')

    return g