from .parser import isint, OPERATORS

def prefix_to_brackets(eqn):
    stack = []
    lastop = []
    intunit = []
    N = len(eqn)
    i = 0
    while i < N:
        # print("Stack", stack)
        val = eqn[i]
        if val.startswith("INT"):
            intunit.append(val)
            i += 1
            while i < N and isint(eqn[i]):
                intunit.append(eqn[i])
                i += 1
            stack.append(" ".join(intunit))
            intunit = []
            i -= 1
        elif val in OPERATORS:
            _, numops = OPERATORS[val]
            lastop.append((len(stack), numops))
            stack.append(val)
        else:
            stack.append(val)

        while len(lastop) > 0  and len(stack) > lastop[-1][0] + lastop[-1][1]:
            # Combine op
            # print(lastop[-1], stack[lastop[-1][0]:])
            op = " ".join(stack[lastop[-1][0]:])
            del stack[lastop[-1][0]:]
            lastop.pop()
            stack.append(f"( {op} )")
        i += 1
    assert(len(stack) == 1)
    return stack[0]

def prefix_to_postfix(eqn):
    if eqn[0].startswith("INT"):
        intunit = [eqn[0]]
        for i, val in enumerate(eqn[1:]):
            if not isint(val):
                break
            intunit.append(val)
        return intunit, eqn[i+1:]
    elif eqn[0] in OPERATORS:
        _, numops = OPERATORS[eqn[0]]
        remeqn = eqn[1:]
        ops = []
        for i in range(numops):
            op, remeqn = prefix_to_postfix(remeqn)
            ops.extend(op)
        ops.append(eqn[0]) # Restructured to postfix
        return ops, remeqn
    else:
        return [eqn[0]], eqn[1:]

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser("Change equation format from prefix to other")
    parser.add_argument("--eqn", required=True)
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    with open(args.eqn, "r") as inf, open(args.out, "w") as outf:
        for eqn in inf:
            postfix, _ = prefix_to_postfix(eqn.strip().split(" "))
            outf.write(" ".join(postfix) + "\n")

    # eqn = "div mul x add INT+ 5 add mul INT+ 3 x mul pow x INT+ 2 add INT- 5 add mul INT- 3 x mul x mul add INT+ 1 mul k0 pow x INT+ 3 add INT+ 4 x add INT+ 5 mul INT+ 3 x"
    # eqn = "div add mul INT+ 3 x pow x INT- 4 mul sub x k0 add mul INT+ 5 x k1"
    # print(" ".join(prefix_to_postfix(eqn.split(" "))[0]))
    # postfix = "x INT+ 3 mul INT+ 5 add x INT+ 4 add INT+ 3 x pow k0 mul INT+ 1 add mul x mul x INT- 3 mul add INT- 5 add INT+ 2 x pow mul x INT+ 3 mul add INT+ 5 add x mul div"
    # print(prefix_to_brackets(eqn.split(" ")))

    # (div (mul x (add INT+ 5 (add (mul INT+ 3 x) (mul (pow x INT+ 2) (add INT- 5 (add (mul INT- 3 x) (mul x (mul (add INT+ 1 (mul k0 (pow x INT+ 3))) (add INT+ 4 x))))))))) (add INT+ 5 (mul INT+ 3 x)))
    # ( div ( mul x ( add INT+ 5 ( add ( mul INT+ 3 x ) ( mul ( pow x INT+ 2 ) ( add INT- 5 ( add ( mul INT- 3 x ) ( mul x ( mul ( add INT+ 1 ( mul k0 ( pow x INT+ 3 ) ) ) ( add INT+ 4 x ) ) ) ) ) ) ) ) ) ( add INT+ 5 ( mul INT+ 3 x ) ) )