<commit_before># Teach Maple (through depends and eval) about our new binding forms.
# Integrand and LO bind from 1st arg to 2nd arg.

`depends/Integrand` := proc(v, e, x) depends(e, x minus {v}) end proc:
`depends/LO`        := proc(v, e, x) depends(e, x minus {v}) end proc:

`eval/Integrand` := proc(e, eqs)
  local v, ee;
  v, ee := op(e);
  eval(op(0,e), eqs)(BindingTools:-generic_evalat(v, ee, eqs))
end proc:

`eval/LO` := proc(e, eqs)
  local v, ee;
  v, ee := op(e);
  eval(op(0,e), eqs)(BindingTools:-generic_evalat(v, ee, eqs))
end proc:

#############################################################################

NewSLO := module ()
  option package;
  local t_pw, t_sum, t_product,
        integrate_known, known_continuous, known_discrete,
        recognize_continuous, recognize_discrete, unroll_GAMMA, get_de, get_se,
        recognize_de, mysolve, Shiftop, Diffop, Recognized,
        factorize, bind, weight,
        reduce_IntSum, reduce_IntsSums, get_indicators,
        elim_intsum, do_elim_intsum, banish, banish_guard, banish_weight,
        reduce_pw, nub_piecewise, piecewise_if,
        find_vars, kb_from_path, interpret, reconstruct, invert,
        get_var_pos, get_int_pos,
        avoid_capture, change_var, disint2,
        mk_sym, mk_ary, mk_idx,
        ModuleLoad;
  export
     # These first few are smart constructors (for themselves):
         integrate, applyintegrand,
     # while these are "proper functions"
         RoundTrip, Simplify, SimplifyKB, TestSimplify, TestHakaru,
         toLO, fromLO, unintegrate, unweight, improve, reduce,
         density, bounds,
         ReparamDetermined, determined, reparam, disint;
  # these names are not assigned (and should not be).  But they are
  # used as global names, so document that here.
  global LO, Integrand, Indicator, SumIE, ProductIE;
  uses Hakaru, KB, Loop;

  RoundTrip := proc(e, t::t_type)
    local result;
    interface(screenwidth=infinity, prettyprint=0);
    writeto("/dev/fd/2");
    result := eval(ToInert(Simplify(_passed)), _Inert_ATTRIBUTE=NULL);
    writeto(terminal);
    lprint(result)
  end proc;

  Simplify := proc(e, t::t_type, {ctx :: list := []}, $)
    evalindets(SimplifyKB(e, t, foldr(assert, empty, op(ctx))),
               And({t_sum, t_product}, anyfunc(anything, anything=range)),
               e -> subsop(0 = `if`(e::t_sum, SumIE, ProductIE),
                           applyop(`+`, [2,2,2], e, 1)))
  end proc;

  SimplifyKB := proc(e, t::t_type, kb::t_kb, $)
    local patterns, x, kb1, ex;
    if t :: HMeasure(anything) then
      fromLO(improve(toLO(e), _ctx=kb), _ctx=kb)
    elif t :: HFunction(anything, anything) then
      patterns := htype_patterns(op(1,t));
      if patterns :: Branches(Branch(PVar(name),anything)) then
        # Eta-expand the function type
        x := `if`(e::lam(name,anything,anything), op(1,e),
                  op([1,1,1],patterns));
        x, kb1 := genType(x, op(1,t), kb, e);
        ex := app(e,x);
        lam(x, op(1,t), SimplifyKB(ex, op(2,t), kb1))
      else
        # Eta-expand the function type and the sum-of-product argument-type
        x := `if`(e::lam(name,anything,anything), op(1,e), d);
        if depends([e,t,kb], x) then x := gensym(x) end if;
        ex := app(e,x);
        lam(x, op(1,t), 'case'(x,
          map(proc(branch)
                local eSubst, pSubst, p1, binds, y, kb1, i, pSubst1;
                eSubst, pSubst := pattern_match([x,e], x, op(1,branch));
                p1 := subs(pSubst, op(1,branch));
                binds := [pattern_binds(p1)];
                kb1 := kb;
                pSubst1 := table();
                for i from 1 to nops(binds) do
                  y, kb1 := genType(op(i,binds), op([2,i],branch), kb1);
                  pSubst1[op(i,binds)] := y;
                end do;
                pSubst1 := op(op(pSubst1));
                Branch(subs(pSubst1, p1),
                       SimplifyKB(eval(eval(ex,eSubst),pSubst1), op(2,t), kb1))
              end proc,
              patterns)))
      end if
    else
      simplify_assuming(e, kb)
    end if
  end proc;

# Testing

  TestSimplify := proc(m, t, n::algebraic:=m, {verify:=simplify})
    local s, r;
    # How to pass keyword argument {ctx::list:=[]} on to Simplify?
    s, r := selectremove(type, [_rest], 'identical(ctx)=anything');
    CodeTools[Test](Simplify(m,t,op(s)), n, measure(verify), op(r))
  end proc;

  TestHakaru := proc(m, n::algebraic:=m,
                     {simp:=improve, verify:=simplify, ctx::list:=[]})
    local kb;
    kb := foldr(assert, empty, op(ctx));
    CodeTools[Test](fromLO(simp(toLO(m), _ctx=kb), _ctx=kb), n,
      measure(verify), _rest)
  end proc;

  t_pw      := 'specfunc(piecewise)';
  t_sum     := 'specfunc({sum    ,Sum    ,%sum    ,%Sum    })';
  t_product := 'specfunc({product,Product,%product,%Product})';

# An integrand h is either an Integrand (our own binding construct for a
# measurable function to be integrated) or something that can be applied
# (probably proc, which should be applied immediately, or a generated symbol).

  applyintegrand := proc(h, x, $)
    local var, body, dummy;
    if h :: 'Integrand(name, anything)' then
      var, body := op(h);
      if x :: {boolean, specfunc({And,Or,Not})} then
        body := eval(body, var=dummy);
        var  := dummy;
        body := subs((var=true)=var, body);
      end if;
      eval(body, var=x)
    elif h :: appliable then
      h(x)
    else
      'procname(_passed)'
    end if
  end proc;

# Step 1 of 3: from Hakaru to Maple LO (linear operator)

  toLO := proc(m, $)
    local h;
    h := gensym('h');
    LO(h, integrate(m, h, []))
  end proc;

  integrate := proc(m, h, loops :: list(name = range) := [], $)
    local x, n, i, res, l;

    if m :: known_continuous then
      integrate_known(Int, Ints, 'xx', m, h, loops)
    elif m :: known_discrete then
      integrate_known(Sum, Sums, 'kk', m, h, loops)
    elif m :: 'Ret(anything)' then
      applyintegrand(h, mk_ary(op(1,m), loops))
    elif m :: 'Bind(anything, name, anything)' then
      res := eval(op(3,m), op(2,m) = mk_idx(op(2,m), loops));
      res := eval(Integrand(op(2,m), 'integrate'(res, x, loops)), x=h);
      integrate(op(1,m), res, loops);
    elif m :: 'specfunc(Msum)' and (nops(loops) = 0 or nops(m) = 1) then
      `+`(op(map(integrate, [op(m)], h, loops)))
    elif m :: 'Weight(anything, anything)' then
      foldl(product, op(1,m), op(loops)) * integrate(op(2,m), h, loops)
    elif m :: t_pw
      and not depends([seq(op(i,m), i=1..nops(m)-1, 2)], map(lhs, loops)) then
      n := nops(m);
      piecewise(seq(`if`(i::even or i=n, integrate(op(i,m), h, loops), op(i,m)),
                    i=1..n))
    elif m :: t_case and not depends(op(1,m), map(lhs, loops)) then
      subsop(2=map(proc(b :: Branch(anything, anything))
                     eval(subsop(2='integrate'(op(2,b), x, loops),b), x=h)
                   end proc,
                   op(2,m)),
             m);
    elif m :: 'LO(name, anything)' then
      eval(op(2,m), op(1,m) = h)
    elif m :: 'Plate(nonnegint, name, anything)' then
      # Unroll Plate when the bound is known. We unroll Plate here (instead
      # of unrolling Ints in reduce, for example) because we have no other
      # way to integrate certain Plates, namely those whose bodies' control
      # flows depend on the index.
      x := mk_sym('pp', h);
      x := [seq(cat(x,i), i=0..op(1,m)-1)];
      if op(1,m) = 0 then
        res := undefined;
      else
        if op(1,m) = 1 then
          res := op(1,x);
        else
           res := idx(x,op(2,m));
#          res := piecewise(seq(op([op(2,m)=i-1, op(i,x)]), i=2..op(1,m)),
#                         op(1,x));
        end if;
        res := mk_idx(res, loops);
      end if;
      res := applyintegrand(h, mk_ary('ary'(op(1,m), op(2,m), res), loops));
      for i from op(1,m) to 1 by -1 do
        res := integrate(eval(op(3,m), op(2,m)=i-1),
                         Integrand(op(i,x), res), loops);
      end do;
      res
    elif m :: 'Plate(anything, name, anything)' then
      integrate(op(3,m), h, [op(2,m)=0..op(1,m)-1, op(loops)])
    elif m :: 'Context(anything, anything)' then
      applyop(integrate, 2, m, h, loops)
    elif h :: appliable then
      x := gensym('xa');
      'integrate'(m, Integrand(x, h(x)), loops)
    else
      'procname(_passed)'
    end if
  end proc;

  integrate_known := proc(make, makes, var, m, h, loops :: list(name=range), $)
    local x, dens, bds;
    x := mk_sym(var, h);
    dens := density[op(0,m)](op(m));
    bds := bounds[op(0,m)](op(m));
    if loops = [] then
      make(dens(x) * applyintegrand(h, x), x = bds);
    else
      makes(foldl(product, dens(mk_idx(x,loops)), op(loops))
              * applyintegrand(h, x),
            x, bds, loops)
    end if;
  end proc;

  known_continuous := '{
    Lebesgue(anything, anything), Uniform(anything, anything),
    Gaussian(anything, anything), Cauchy(anything, anything),
    StudentT(anything, anything, anything), ChiSquared(anything),
    BetaD(anything, anything), GammaD(anything, anything)}':

  known_discrete := '{Counting(anything, anything),
    Categorical(anything),
    NegativeBinomial(anything), PoissonD(anything)}';<commit_msg>Corrected number of args of NegativeBinomail to 2.<commit_after># Teach Maple (through depends and eval) about our new binding forms.
# Integrand and LO bind from 1st arg to 2nd arg.

`depends/Integrand` := proc(v, e, x) depends(e, x minus {v}) end proc:
`depends/LO`        := proc(v, e, x) depends(e, x minus {v}) end proc:

`eval/Integrand` := proc(e, eqs)
  local v, ee;
  v, ee := op(e);
  eval(op(0,e), eqs)(BindingTools:-generic_evalat(v, ee, eqs))
end proc:

`eval/LO` := proc(e, eqs)
  local v, ee;
  v, ee := op(e);
  eval(op(0,e), eqs)(BindingTools:-generic_evalat(v, ee, eqs))
end proc:

#############################################################################

NewSLO := module ()
  option package;
  local t_pw, t_sum, t_product,
        integrate_known, known_continuous, known_discrete,
        recognize_continuous, recognize_discrete, unroll_GAMMA, get_de, get_se,
        recognize_de, mysolve, Shiftop, Diffop, Recognized,
        factorize, bind, weight,
        reduce_IntSum, reduce_IntsSums, get_indicators,
        elim_intsum, do_elim_intsum, banish, banish_guard, banish_weight,
        reduce_pw, nub_piecewise, piecewise_if,
        find_vars, kb_from_path, interpret, reconstruct, invert,
        get_var_pos, get_int_pos,
        avoid_capture, change_var, disint2,
        mk_sym, mk_ary, mk_idx,
        ModuleLoad;
  export
     # These first few are smart constructors (for themselves):
         integrate, applyintegrand,
     # while these are "proper functions"
         RoundTrip, Simplify, SimplifyKB, TestSimplify, TestHakaru,
         toLO, fromLO, unintegrate, unweight, improve, reduce,
         density, bounds,
         ReparamDetermined, determined, reparam, disint;
  # these names are not assigned (and should not be).  But they are
  # used as global names, so document that here.
  global LO, Integrand, Indicator, SumIE, ProductIE;
  uses Hakaru, KB, Loop;

  RoundTrip := proc(e, t::t_type)
    local result;
    interface(screenwidth=infinity, prettyprint=0);
    writeto("/dev/fd/2");
    result := eval(ToInert(Simplify(_passed)), _Inert_ATTRIBUTE=NULL);
    writeto(terminal);
    lprint(result)
  end proc;

  Simplify := proc(e, t::t_type, {ctx :: list := []}, $)
    evalindets(SimplifyKB(e, t, foldr(assert, empty, op(ctx))),
               And({t_sum, t_product}, anyfunc(anything, anything=range)),
               e -> subsop(0 = `if`(e::t_sum, SumIE, ProductIE),
                           applyop(`+`, [2,2,2], e, 1)))
  end proc;

  SimplifyKB := proc(e, t::t_type, kb::t_kb, $)
    local patterns, x, kb1, ex;
    if t :: HMeasure(anything) then
      fromLO(improve(toLO(e), _ctx=kb), _ctx=kb)
    elif t :: HFunction(anything, anything) then
      patterns := htype_patterns(op(1,t));
      if patterns :: Branches(Branch(PVar(name),anything)) then
        # Eta-expand the function type
        x := `if`(e::lam(name,anything,anything), op(1,e),
                  op([1,1,1],patterns));
        x, kb1 := genType(x, op(1,t), kb, e);
        ex := app(e,x);
        lam(x, op(1,t), SimplifyKB(ex, op(2,t), kb1))
      else
        # Eta-expand the function type and the sum-of-product argument-type
        x := `if`(e::lam(name,anything,anything), op(1,e), d);
        if depends([e,t,kb], x) then x := gensym(x) end if;
        ex := app(e,x);
        lam(x, op(1,t), 'case'(x,
          map(proc(branch)
                local eSubst, pSubst, p1, binds, y, kb1, i, pSubst1;
                eSubst, pSubst := pattern_match([x,e], x, op(1,branch));
                p1 := subs(pSubst, op(1,branch));
                binds := [pattern_binds(p1)];
                kb1 := kb;
                pSubst1 := table();
                for i from 1 to nops(binds) do
                  y, kb1 := genType(op(i,binds), op([2,i],branch), kb1);
                  pSubst1[op(i,binds)] := y;
                end do;
                pSubst1 := op(op(pSubst1));
                Branch(subs(pSubst1, p1),
                       SimplifyKB(eval(eval(ex,eSubst),pSubst1), op(2,t), kb1))
              end proc,
              patterns)))
      end if
    else
      simplify_assuming(e, kb)
    end if
  end proc;

# Testing

  TestSimplify := proc(m, t, n::algebraic:=m, {verify:=simplify})
    local s, r;
    # How to pass keyword argument {ctx::list:=[]} on to Simplify?
    s, r := selectremove(type, [_rest], 'identical(ctx)=anything');
    CodeTools[Test](Simplify(m,t,op(s)), n, measure(verify), op(r))
  end proc;

  TestHakaru := proc(m, n::algebraic:=m,
                     {simp:=improve, verify:=simplify, ctx::list:=[]})
    local kb;
    kb := foldr(assert, empty, op(ctx));
    CodeTools[Test](fromLO(simp(toLO(m), _ctx=kb), _ctx=kb), n,
      measure(verify), _rest)
  end proc;

  t_pw      := 'specfunc(piecewise)';
  t_sum     := 'specfunc({sum    ,Sum    ,%sum    ,%Sum    })';
  t_product := 'specfunc({product,Product,%product,%Product})';

# An integrand h is either an Integrand (our own binding construct for a
# measurable function to be integrated) or something that can be applied
# (probably proc, which should be applied immediately, or a generated symbol).

  applyintegrand := proc(h, x, $)
    local var, body, dummy;
    if h :: 'Integrand(name, anything)' then
      var, body := op(h);
      if x :: {boolean, specfunc({And,Or,Not})} then
        body := eval(body, var=dummy);
        var  := dummy;
        body := subs((var=true)=var, body);
      end if;
      eval(body, var=x)
    elif h :: appliable then
      h(x)
    else
      'procname(_passed)'
    end if
  end proc;

# Step 1 of 3: from Hakaru to Maple LO (linear operator)

  toLO := proc(m, $)
    local h;
    h := gensym('h');
    LO(h, integrate(m, h, []))
  end proc;

  integrate := proc(m, h, loops :: list(name = range) := [], $)
    local x, n, i, res, l;

    if m :: known_continuous then
      integrate_known(Int, Ints, 'xx', m, h, loops)
    elif m :: known_discrete then
      integrate_known(Sum, Sums, 'kk', m, h, loops)
    elif m :: 'Ret(anything)' then
      applyintegrand(h, mk_ary(op(1,m), loops))
    elif m :: 'Bind(anything, name, anything)' then
      res := eval(op(3,m), op(2,m) = mk_idx(op(2,m), loops));
      res := eval(Integrand(op(2,m), 'integrate'(res, x, loops)), x=h);
      integrate(op(1,m), res, loops);
    elif m :: 'specfunc(Msum)' and (nops(loops) = 0 or nops(m) = 1) then
      `+`(op(map(integrate, [op(m)], h, loops)))
    elif m :: 'Weight(anything, anything)' then
      foldl(product, op(1,m), op(loops)) * integrate(op(2,m), h, loops)
    elif m :: t_pw
      and not depends([seq(op(i,m), i=1..nops(m)-1, 2)], map(lhs, loops)) then
      n := nops(m);
      piecewise(seq(`if`(i::even or i=n, integrate(op(i,m), h, loops), op(i,m)),
                    i=1..n))
    elif m :: t_case and not depends(op(1,m), map(lhs, loops)) then
      subsop(2=map(proc(b :: Branch(anything, anything))
                     eval(subsop(2='integrate'(op(2,b), x, loops),b), x=h)
                   end proc,
                   op(2,m)),
             m);
    elif m :: 'LO(name, anything)' then
      eval(op(2,m), op(1,m) = h)
    elif m :: 'Plate(nonnegint, name, anything)' then
      # Unroll Plate when the bound is known. We unroll Plate here (instead
      # of unrolling Ints in reduce, for example) because we have no other
      # way to integrate certain Plates, namely those whose bodies' control
      # flows depend on the index.
      x := mk_sym('pp', h);
      x := [seq(cat(x,i), i=0..op(1,m)-1)];
      if op(1,m) = 0 then
        res := undefined;
      else
        if op(1,m) = 1 then
          res := op(1,x);
        else
           res := idx(x,op(2,m));
#          res := piecewise(seq(op([op(2,m)=i-1, op(i,x)]), i=2..op(1,m)),
#                         op(1,x));
        end if;
        res := mk_idx(res, loops);
      end if;
      res := applyintegrand(h, mk_ary('ary'(op(1,m), op(2,m), res), loops));
      for i from op(1,m) to 1 by -1 do
        res := integrate(eval(op(3,m), op(2,m)=i-1),
                         Integrand(op(i,x), res), loops);
      end do;
      res
    elif m :: 'Plate(anything, name, anything)' then
      integrate(op(3,m), h, [op(2,m)=0..op(1,m)-1, op(loops)])
    elif m :: 'Context(anything, anything)' then
      applyop(integrate, 2, m, h, loops)
    elif h :: appliable then
      x := gensym('xa');
      'integrate'(m, Integrand(x, h(x)), loops)
    else
      'procname(_passed)'
    end if
  end proc;

  integrate_known := proc(make, makes, var, m, h, loops :: list(name=range), $)
    local x, dens, bds;
    x := mk_sym(var, h);
    dens := density[op(0,m)](op(m));
    bds := bounds[op(0,m)](op(m));
    if loops = [] then
      make(dens(x) * applyintegrand(h, x), x = bds);
    else
      makes(foldl(product, dens(mk_idx(x,loops)), op(loops))
              * applyintegrand(h, x),
            x, bds, loops)
    end if;
  end proc;

  known_continuous := '{
    Lebesgue(anything, anything), Uniform(anything, anything),
    Gaussian(anything, anything), Cauchy(anything, anything),
    StudentT(anything, anything, anything), ChiSquared(anything),
    BetaD(anything, anything), GammaD(anything, anything)}':

  known_discrete := '{Counting(anything, anything),
    Categorical(anything),
    NegativeBinomial(anything,anything), PoissonD(anything)}';
{'fraction_non_alphanumeric': 0.16815492200107585, 'fraction_numerical': 0.010704679935449166, 'mean_word_length': 2.7189437887577514, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1825587', 'n_tokens_mistral': 6862, 'n_tokens_neox': 6533, 'n_words': 1706}