use symbolica::prelude::*;

fn main() {
    let input = parse!("x*(1+a)+x*5*y+f(5,x)+2+y^2+x^2 + x^3");
    let x = Atom::var(symbol!("x"));
    let key = symbol!("key");
    let coeff = symbol!("val");

    let r = input.coefficient_list::<i8>(std::slice::from_ref(&x));

    println!("> Coefficient list:");
    for (key, val) in r {
        println!("\t{key} {val}");
    }

    println!("> Collect in x:");
    let out = input.collect::<i8>(&x);
    println!("\t{out}");

    println!("> Collect in x with wrapping:");
    let out = input.collect_mapped::<i8>(
        &x,
        move |a, out| {
            out.set_from_view(&a);
            **out = function!(key, out.as_view());
        },
        move |a, out| {
            out.set_from_view(&a);
            **out = function!(coeff, out.as_view());
        },
    );
    println!("\t{out}");
}
