/* Term generation */
exp_taylor(expr, n) := expand(subst([ξ=expr], taylor(exp(ξ), ξ, 0, n)));
neg_pow_series(c, r, a, b) := sum(c[i] / r^i, i, a, b);


/* Term filtering */

filter_order(expr, x, a, b) := block([acc, len],
  acc : 0,
  len : if atom(expr) or not(op(expr) = "+") then 1 else length(expr),

  for i thru len do block([deg, p],
    p : if len=1 then expr else part(expr, i),
    deg : hipow(partition([p], x)[2], x),
    if (deg >= a and deg <= b) then acc : acc + p),

  acc);

/* Speed up by expanding and mapping only once over the list (check operator is +) */
keep_upto_order(expr, var, start, end) := block([parts],
  expr : expand(expr),
  parts : makelist(coeff(expr, var, i) * r^i, i, start, end),
  apply("+", parts));


/* Term generation with filtering */

/*
Calculate the exponential series expansion evaluated at `expr' up to `n'th order
in `x'
*/
exp_series(expr, x, n, a, b) := block([acc, xn, len],
  expr : expand(expr),
  xn : 1,
  acc : 1,
  len : if atom(expr) or not(is(op(expr)="+")) then 1 else length(expr),
  /* Assume a, b < 0 for now */
  for i thru n do block([subexpr],
    /* Discard terms that would end up creating terms outside range */
    /*
    subexpr : 0,
    for j thru len do block([term, deg],
      term : if len=1 then expr else part(expr, j),
      deg : hipow(term, x),
      if not((deg + lopow(xn, x)) > b or (deg + hipow(xn, x)) < a)
      then
      subexpr : subexpr + term),
    */
    /**/
    xn : filter_order(ratexpand(xn * expr / i), x, a, b),
    acc : acc + xn),
  acc);

horner_exp_series(expr, x, n, a, b) := block([acc],
  /* expand(1 + x*(1 + x*(1 + x*(1+x/4)/3)/2)) */
  expr : expand(expr),
  acc : 1,
  for i thru n do acc : keep_upto_order(1 + expr*acc/(n-i+1), x, a, b),
  acc);


/* Multiply series expansions upto nth order */

matchdeclare(m_freeof_r, lambda([e], freeof(r, e)))$
matchdeclare(m_small, lambda([e], e < -expansion_order))$
defrule(r_simp_small, m_freeof_r*r^m_small, 0)$

__simple_filter(series, max_order) := block([ret, tmp],
  tmp : expansion_order,
  expansion_order : max_order,
  ret : apply1(series, r_simp_small),
  expansion_order : tmp,
  ret);

__table_size(factors) := block(
  [nfactors, ranges, t_min, t_max],
  nfactors : length(factors),
  ranges : make_array(any, nfactors, 2),
  /* Fill table with maximal and minimal powers */
  for i : 0 thru nfactors-1 do block(
    ranges[i, 0] : lopow(factors[i+1], r),
    ranges[i, 1] : hipow(factors[i+1], r)),
  /* Find greatest and smallest r powers */
  t_min : ranges[0, 0],
  t_max : ranges[0, 1],
  for i : 0 thru nfactors-1 do (
    t_min : min(t_min, ranges[i, 0]),
    t_max : max(t_max, ranges[i, 1])),
  [t_min, t_max]);

__fill_term_table(table, factors, t_min, t_max, len) := block(
  [nfactors],
  nfactors : length(factors),
  /* Split factors into term table */
  for i : 0 thru nfactors-1 do (
    /* Initialize table row with 0s */
    for j : 0 thru len-1 do table[i, j] : 0,
    /**/
    if not atom(factors[i+1]) and op(factors[i+1])="+" then (
      for term in args(factors[i+1]) do block(
        [j, n],
        n : lopow(term, r),
        j : n - t_min,
        table[i, j] : table[i, j] + coeff(term, r, n)))
    else block(
      [j, n, term],
      term : factors[i+1],
      n : lopow(term, r),
      j : n  - t_min,
      table[i, j] : coeff(term, r, n))));

__mul_series_level(table, max_order, i, stack, nfactors, len, t_min) := block(
  [],
  if is(i = nfactors) then block(
    [order, acc],
    order : 0,
    for j : 0 thru nfactors-1 do order : order + stack[j] + t_min,
    if order >= -max_order then (
      acc : 1,
      for j : 0 thru nfactors-1 do acc : acc * table[j, stack[j]],
      acc * r^order)
    else 0)
  else block([acc],
    /* TODO: omit external 0s */
    acc : 0,
    for j : 0 thru len-1 do (
      stack[i] : j,
      acc : acc + __mul_series_level(
        table, max_order, i+1, stack, nfactors, len, t_min)),
    acc));

/*
Multiply list of factors, each expressed as terms of series, only upto max_order
*/
mul_series(factors, max_order) := block(
  [term_table, len, t_min, t_max, stack, nfactors],
  nfactors : length(factors),

  /* If there is nothing to multiple, still filter higher terms */
  if is(nfactors=1) then (
    return(__simple_filter(factors[1], max_order))),

  /* Allocate stack and table for ranges */
  stack : make_array(any, nfactors),

  /* Compute table size */
  block([pair],
    pair : __table_size(factors),
    t_min : pair[1],
    t_max : pair[2]),
  len : abs(t_max - t_min) + 1,

  /* Allocate table for terms */
  term_table : make_array(any, nfactors, len),
  __fill_term_table(term_table, factors, t_min, t_max, len),

  __mul_series_level(term_table, max_order, 0, stack, nfactors, len, t_min))$


/**/

β_expr_val(x, n) := block(
  if x < 0 then (
    exp(x*H)*horner_exp_series(abs(x)*(H-rhs(main_ansatz[2])), r, n, -n, 2))
  else (
    exp(x*H)*horner_exp_series(x*(rhs(main_ansatz[2])-H), r, n, -n, 2)));

γ_expr_val(x, n) := horner_exp_series(x*rhs(main_ansatz[1]), r, n, -n, 2);

matchdeclare(m_any1, lambda ([e], freeof(β, γ, e)));
matchdeclare(m_any2, lambda ([e], freeof(β, γ, e)));
matchdeclare(m_any3, lambda ([e], freeof(β, γ, e)));
matchdeclare(m_beta, lambda ([e], not freeof(β, e) and freeof(γ, e)));
matchdeclare(m_gamma, lambda ([e], not freeof(γ, e) and freeof(β, e)));

/* Replace β and γ expressions into a symbolic representation */
defrule(r_wrap_exp_β, exp(m_any1*m_beta), β_expr(m_any1));
defrule(r_wrap_exp_γ, exp(m_any1*m_gamma), γ_expr(m_any1));
defrule(r_wrap_exp_βγ, exp(m_any1*m_beta + m_any2*m_gamma),
  β_expr(m_any1)*γ_expr(m_any2));
defrule(r_wrap_exp_βγ2, exp(m_any1*m_beta + m_any2*m_gamma + m_any3),
  exp(m_any3)*β_expr(m_any1)*γ_expr(m_any2));

/* Apply wrap rules to expression */
wrap_exp(expr) := (
  apply1(expr, r_wrap_exp_β, r_wrap_exp_γ, r_wrap_exp_βγ, r_wrap_exp_βγ2));

/* Replace symbolic representations of β and γ expressions into their their
expansions */
defrule(r_exp_β_val, β_expr(m_any1), β_expr_val(m_any1, expansion_order));
defrule(r_exp_γ_val, γ_expr(m_any1), γ_expr_val(m_any1, expansion_order));

unwrap_exp(expr) := apply1(expr, r_exp_β_val, r_exp_γ_val);


/**/

freeof_many(expr, funcs) := apply(freeof, append(funcs, [expr]));

/* Split product in factors containing any funcs and factors free of them */
split_prod(expr, funcs) := block(
  [deps_on_funcs : []],
  /* If not a product, assume product with single factor */
  for f in funcs do block(
    if atom(expr) then (
      if member(expr, funcs) then (
        push(expr, deps_on_funcs),
        expr : 1),
      return())
    elseif op(expr)="*" then block(
      [split : partition(expr, f)],
      expr : split[1],
      push(split[2], deps_on_funcs))
    elseif not freeof_many(expr, funcs) then (
      push(expr, deps_on_funcs),
      expr : 1,
      return())
    else (
      return())),
  [expr, apply("*", deps_on_funcs)]);

/* Substitute functions by their series expansions and apply substitutions in
coeffs */
__expand_term_subst(expr, coeffs) := block(
  expr : apply1(expr, r_exp_β_val, r_exp_γ_val),
  expr : subst(append(main_ansatz, coeffs), expr),
  ev(expr, eval, diff, expand));

/* Substitute product of functions by result of mul_series */
__expand_term_prod(expr, coeffs) := block(
  [free, factors],
  expr : wrap_exp(expr),
  /* Split product into list of factors containing funcs and those independent
  of funcs*/
  block([split : split_prod(expr, funcs)],
    free : split[1],
    if not atom(split[2]) and op(split[2])="*" then (
      factors : args(split[2]))
    else (
      factors : [split[2]])),
  /* Substitute funcs factors by expansion */
  factors : __expand_term_subst(factors, coeffs),
  /* Total product using series optimization */
  free * mul_series(factors, expansion_order + lopow(free, r)));

/*
Split product factors into list substitute series expansions, substitute
coefficients and finally multiply only upto expansion_order.
*/
expand_term(expr, coeffs) := block(
  [tmp],
  if atom(expr) then (
    apply1(__expand_term_subst(wrap_exp(expr), coeffs), r_simp_small))
  /* TODO: Recursive call instead of error */
  elseif op(expr) = "+" then (
    error("Term is itself a sum!"))
  /* For single negation just call recursively */
  elseif op(expr) = "-" then (
    if length(args(expr)) # 1 then (
      error("Unexpected '-' operator with more than 1 argument")),
    expand(- expand_term(part(expr, 1), coeffs)))
  elseif op(expr) = "*" or op(expr) = "^" then (
    expand(__expand_term_prod(expr, coeffs)))
  elseif op(expr) = "/" then block(
    [n : part(expr, 1), d : part(expr, 2), expansion_order : expansion_order],
    expansion_order : expansion_order - lopow(d, r),
    expand(expand_term(n, coeffs) / d))
  else (
    error(printf(false, "expand_term: invalid operator '~a'", op(expr)))));


/* Solve main equations term by term */

print_progress : false;

/* TODO: Divide and conquer on term list for long expansions */
/* coeffs: facts already known about series expansions */

/* Series expansion inside each instance of a value in the terms of an equation
*/
expand_equation(expr, coeffs) := (
  expr : rhs(expr) - lhs(expr),
  0 = if atom(expr) then (
    expr)
  elseif op(expr)="+" then block(
    [i : 0, len : length(args(expr))],
    apply("+", map(lambda([term],
          i : i+1,
          if print_progress then (
            printf(true, "expand_equation [~a/~a]~%", i, len)),
          expand_term(term, coeffs)),
        args(expr))))
  else (
    expand_term(expr, coeffs)))$


/* Misc */

/* For an equation of polynomials, returns the term by term list of equations of
the coefficients */
coeff_equations(eq) := block([start, end],
  eq : expand(eq),
  start : lopow(eq, r),
  end : hipow(eq, r),
  reverse(makelist([r^i, coeff(eq, r, i)], i, start, end)));