import sympy as sp
from functools import reduce

def taint_leaves(expr, taintp):
    tainted = set()
    stack = [expr]
    stack_iters = [iter(expr.args)]
    while True:
        while stack and (n := next(stack_iters[-1], None)) is None:
            stack.pop()
            stack_iters.pop()
        if not stack:
            break

        if taintp(n):
            for p in stack:
                tainted.add(id(p))
                tainted.add(id(n))

        if n.args:
            stack.append(n)
            stack_iters.append(iter(n.args))

    return tainted

def find_nodes(expr, taintp):
    tainted = set()
    if taintp(expr):
        tainted.add(id(expr))
    stack = [expr]
    stack_iters = [iter(expr.args)]
    while True:
        while stack and (n := next(stack_iters[-1], None)) is None:
            stack.pop()
            stack_iters.pop()
        if not stack:
            break

        if n.args:
            stack.append(n)
            stack_iters.append(iter(n.args))

        if taintp(n):
            tainted.add(id(n))

    return tainted

def _combine_specials_list(a, b):
    return [
        (ia[0] * ib[0],) + ia[1:] + ib[1:]
        for ia in a
        for ib in b
    ]

def _single_out_special_terms(expr, descend, specials):
    if id(expr) not in descend:
        return [(expr,)]

    if id(expr) in specials:
        return [(sp.S.One, expr)]
    
    if type(expr) is sp.Add:
        return [term for a in expr.args
            for term in _single_out_special_terms(
                a, descend, specials
            )]

    if type(expr) is sp.Mul:
        return reduce(
            _combine_specials_list,
            [_single_out_special_terms(i, descend, specials)
                for i in expr.args]
        )

    if type(expr) is sp.Pow and type(expr.args[1]) is sp.Integer \
            and expr.args[1] > 0:

        return reduce(
            _combine_specials_list,
            [_single_out_special_terms(expr.args[0], descend, specials)] * int(expr.args[1])
        )

    raise NotImplementedError(expr)

def _collect_special_terms(list_):
    ret = dict()
    for decomp in list_:
        poly, special = decomp[0], decomp[1:]
        if special not in ret:
            ret[special] = poly
        else:
            ret[special] += poly
    return [
        (value,) + key
        for key, value in ret.items()
        if value != sp.S.Zero
    ]

def single_out_special_terms(expr, var):
    all_specials = find_nodes(expr, lambda v: type(v) not in [sp.Mul, sp.Add] \
        and not isinstance(v, sp.Symbol) and not isinstance(v, sp.Number) and \
        (type(v) is not sp.Pow or type(v.args[1]) is not sp.Integer or v.args[1] < 0))
    specials = taint_leaves(expr, taintp=lambda v: v == var) \
                .intersection(all_specials)
    descend_this_way = taint_leaves(expr, taintp=lambda v: id(v) in specials)
    return _collect_special_terms(_single_out_special_terms(expr, descend_this_way, specials))

def product(*args):
    ret = sp.S.One
    for a in args:
        ret *= a
    return ret

def reform_wrt_specials(expr, var, passes=[sp.cancel]):
    ret = sp.S.Zero
    for comp in single_out_special_terms(expr, var):
        poly, other = comp[0], product(*comp[1:])
        ret += sp.cancel(poly)*other
    return ret

def analyze(expr, var):
    for compo in single_out_special_terms(expr, var):
        print(*compo[1:], compo[0])

def analyze_table(expr, var):
    return sp.Array([
        [product(*compo[1:]), compo[0]]
        for compo in single_out_special_terms(expr, var)
    ])

def attempt_poly_reduce(expr, s):
    ret = sp.S.Zero
    for d in single_out_special_terms(expr, s):
        powers = dict()
        poly = d[0]
        for spec in d[1:]:
            if not spec in powers:
                powers[spec] = 0
            powers[spec] += 1
            if sp.cancel(sp.expand_func(spec**powers[spec]*s**2)).is_polynomial(s):
                poly *= sp.expand_func(spec**powers[spec])
                powers[spec] = 0
        ret += reduce(lambda a, b: a*b, [spec**pow for (spec, pow) in powers.items()], sp.S.One)*poly
    return ret
