import sympy as sp
import random

from .parser import parse_prefix_to_sympy

isconst = lambda e: not any(c.is_symbol for c in e.atoms())
def constfold(expr):
    q = [expr]
    cidx = 0
    subsmap = {}
    constmap = {}

    while len(q) > 0:
        curr_expr = q.pop(0)
        if isinstance(curr_expr, sp.Number) or isconst(e):
            const_expr = curr_expr.evalf()
            rep = sp.nsimplify(const_expr, [sp.E, sp.pi])
            if isinstance(rep, sp.Integer) or \
                    (isinstance(rep, sp.Rational) and rep.q <= 16):
                subsmap[curr_expr] = rep
            else:
                subsmap[curr_expr] = sp.Symbol(f"k{cidx}")
                constmap[f"k{cidx}"] = float(const_expr)
                cidx += 1
        else:
            for child in curr_expr.args:
                q.append(child)

    return expr.subs(subsmap), constmap

def replace_const(expr):
    cidx = 0
    subsmap = {}
    constmap = {}

    for c in sp.preorder_traversal(expr):
        if isinstance(c, sp.Float):
            rep = sp.nsimplify(c)
            if isinstance(rep, sp.Integer) or \
                    (isinstance(rep, sp.Rational) and rep.q <= 16):
                subsmap[c] = rep
            else:
                subsmap[c] = sp.Symbol(f"c{cidx}")
                constmap[f"c{cidx}"] = float(c)
                cidx += 1
    return expr.subs(subsmap), constmap


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser("Random experiments")
    parser.add_argument("-f", required=True)
    parser.add_argument("-p", type=float, default=0.1)
    parser.add_argument("-n", type=int, default=20)
    args = parser.parse_args()

    random.seed(1225)
    
    count = 0
    with open(args.f, "r") as f:
        for line in f:
            if random.random() > args.p:
                continue

            prefl = line.strip().split(" ")

            orig = parse_prefix_to_sympy(prefl)
            # simp = sp.simplify(expr)
            expr = constfold(orig)
            expr, consts = replace_const(expr)
            print(orig, expr, consts)

            count += 1
            if count == args.n:
                break