Raise expectations in gmm (Gaussian mixture model) test

We now recognize the conjugacy and produce index manipulations.
The test fails because the test doesn't allow for alpha-renaming.

kernelopts(assertlevel=2): # be strict on all assertions while testing
kernelopts(opaquemodules=false): # allow testing of internal routines
if not (NewSLO :: `module`) then
  WARNING("loading NewSLO failed");
  `quit`(3);
end if;

with(Hakaru):
with(NewSLO):

#####################################################################
#
# plate/array tests
#
#####################################################################

triv := Plate(k, i, Ret(i)):
TestHakaru(triv, Ret(ary(k, i, i)), label="Dirac Plate");

ary0 := Bind(Plate(k, i, Gaussian(0,1)), xs, Ret([xs])):
TestHakaru(ary0, ary0, label="plate roundtripping", ctx = KB:-assert(k::nonnegint,KB:-empty));

ary1  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(Unit))), ys,
         Ret(x))):
ary1w := 2^(-(1/2)*n+1/2)*exp((1/2)*((sum(idx(t,i),i=0..n-1))^2-(sum(idx(t,i)^2,i=0..n-1))*n-(sum(idx(t,i)^2,i=0..n-1)))/(n+1))*Pi^(-(1/2)*n)/sqrt(2+2*n):
TestHakaru(ary1,
  Weight(ary1w, Gaussian((sum(idx(t, i), i = 0 .. n-1))/(n+1), 1/sqrt(n+1))),
  label="Wednesday goal", ctx = KB:-assert(n::nonnegint,KB:-empty));
TestHakaru(Bind(ary1, x, Ret(Unit)), Weight(ary1w, Ret(Unit)),
  label="Wednesday goal total", ctx = KB:-assert(n::nonnegint,KB:-empty));
ary2  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Bind(Gaussian(idx(t,i),1),z, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(z+1)))), ys,
         Ret(ys))):
TestHakaru(ary2,
  Weight(ary1w, Bind(Plate(n, i, Gaussian(idx(t,i),1)), zs, Ret(ary(n, i, idx(zs,i)+1)))),
  label="Reason for fission", ctx = KB:-assert(n::nonnegint,KB:-empty));
ary3  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Bind(Gaussian(idx(t,i),1),z, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(z)))), zs,
         Ret(zs))):
TestHakaru(ary3, Weight(ary1w, Plate(n, i, Gaussian(idx(t,i),1))),
  label="Array eta", ctx = KB:-assert(n::nonnegint,KB:-empty));
TestHakaru(Ret(ary(n,i,idx(f(i),i))), label="Don't overdo array eta");

bry1  := Bind(BetaD(alpha,beta), x,
         Bind(Plate(n, i, Weight(x    ^piecewise(idx(y,i)=true ,1) *
                                 (1-x)^piecewise(idx(y,i)=false,1),
                          Ret(Unit))), ys,
         Ret(x))):
bry1s := Weight(Beta(alpha+sum(piecewise(idx(y,i)=true ,1), i=0..n-1),
                     beta +sum(piecewise(idx(y,i)=false,1), i=0..n-1))/Beta(alpha,beta),
         BetaD(alpha+sum(piecewise(idx(y,i)=true ,1), i=0..n-1),
               beta +sum(piecewise(idx(y,i)=false,1), i=0..n-1))):
TestHakaru(bry1, bry1s,
  label="first way to express flipping a biased coin many times",
  ctx = KB:-assert(n::nonnegint,KB:-empty));

bry2  := Bind(BetaD(alpha,beta), x,
         Bind(Plate(n, i, Weight(x    ^(  idx(y,i)) *
                                 (1-x)^(1-idx(y,i)),
                          Ret(Unit))), ys,
         Ret(x))):
bry2s := Weight(Beta(alpha+  sum(idx(y,i),i=0..n-1),
                     beta +n-sum(idx(y,i),i=0..n-1))/Beta(alpha,beta),
         BetaD(alpha+  sum(idx(y,i),i=0..n-1),
               beta +n-sum(idx(y,i),i=0..n-1))):
TestHakaru(bry2, bry2s,
  label="second way to express flipping a biased coin many times",
  ctx = KB:-assert(n::nonnegint,KB:-empty));

fission     := Bind(Plate(k, i, Gaussian(z^i,1)), xs, Plate(k, i, Gaussian(idx(xs,i),1))):
fusion      := Plate(k, i, Bind(Gaussian(z^i,1), x, Gaussian(x,1))):
conjugacies := Plate(k, i, Gaussian(z^i, sqrt(2))):
conjugacies5:= Bind(Gaussian(z^0, sqrt(2)), x0,
               Bind(Gaussian(z^1, sqrt(2)), x1,
               Bind(Gaussian(z^2, sqrt(2)), x2,
               Bind(Gaussian(z^3, sqrt(2)), x3,
               Bind(Gaussian(z^4, sqrt(2)), x4,
               Ret(ary(5, i, piecewise(i=1,x1, i=2,x2, i=3,x3, i=4,x4, x0)))))))):
TestHakaru(eval(fission,{    z=1}), eval(conjugacies ,z=1), verify=normal, label="Conjugacy across iid plates");
TestHakaru(eval(fusion ,{    z=1}), eval(conjugacies ,z=1), verify=normal, label="Conjugacy in iid plate");
TestHakaru(eval(fission,{k=5,z=1}), eval(conjugacies5,z=1), verify=normal, label="Conjugacy across iid plates unrolled");
TestHakaru(eval(fusion ,{k=5,z=1}), eval(conjugacies5,z=1), verify=normal, label="Conjugacy in iid plate unrolled");
TestHakaru(     fission           ,      conjugacies      , verify=normal, label="Conjugacy across plates", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(     fusion            ,      conjugacies      , verify=normal, label="Conjugacy in plate", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(eval(fission,{k=5    }),      conjugacies5     , verify=normal, label="Conjugacy across plates unrolled", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(eval(fusion ,{k=5    }),      conjugacies5     , verify=normal, label="Conjugacy in plate unrolled", ctx=KB:-assert(z>0,KB:-empty));

# Simplifying gmm below is a baby step towards index manipulations we need
gmm := Bind(Plate(k, c, Gaussian(0,1)), xs,
       Bind(Plate(n, i, Weight(density[Gaussian](idx(xs,idx(cs,i)),1)(idx(t,i)), Ret(Unit))), ys,
       Ret(xs))):
gmm_s := Weight(2^(-(1/2)*n)*Pi^(-(1/2)*n)*exp(-(1/2)*(sum(idx(t,i)^2, i=0..n-1)))*exp((1/2)*(sum((sum(piecewise(c=idx(cs,i), idx(t,i), 0), i=0..n-1))^2/(sum(piecewise(c=idx(cs,i), 1, 0), i=0..n-1)+1), c=0..k-1)))*(1/(product(sum(piecewise(c=idx(cs,i), 1, 0), i=0..n-1)+1, c=0..k-1)))^(1/2),
         Plate(k, c, Gaussian((sum(piecewise(c=idx(cs,i), idx(t,i), 0), i=0..n-1))/(sum(piecewise(c=idx(cs,i), 1, 0), i=0..n-1)+1), (1/(sum(piecewise(c=idx(cs,i), 1, 0), i=0..n-1)+1))^(1/2)))):
TestHakaru(gmm, gmm_s, label="gmm (currently failing -- due to alpha-equivalence)");

# Detecting Dirichlet-multinomial conjugacy when unrolled
dirichlet := proc(as)
  Bind(Plate(size(as)-1, i, BetaD(sum(idx(as,j), j=0..i), idx(as,i+1))), xs,
  Ret(ary(size(as), i,
          product(idx(xs,j), j=i..size(as)-2)
          * piecewise(0=i, 1, 1-idx(xs,i-1)))))
end proc:
categorical := proc(ps)
  Bind(Counting(0, size(ps)-1), i, Weight(idx(ps,i), Ret(i)))
end proc:
# The next line eta-expands CodeTools[Test] so that its arguments get evaluated
(proc() CodeTools[Test](_passed) end proc)
  (op(2, [unweight(
     fromLO(improve(toLO(
       Plate(n,k,
         Bind(dirichlet(ary(5,i,1)),ps,
           Weight(product(idx(ps,i)^idx(idx(t,k),i),i=0..size(ps)-1),
             Ret(ps))))))))]),
   fromLO(toLO(
     Plate(n,k,
       dirichlet(ary(5,i,1+idx(idx(t,k),i)))))),
   measure(simplify),
   label="Conjugacy between unrolled symmetric Dirichlet and multinomial");
# We'd like the test above to pass even if the count 5 becomes symbolic.
# Below is some progress towards this goal:
TestHakaru(dirichlet(as), label="Dirichlet symbolic prior roundtrip");
TestHakaru(Bind(dirichlet(as),ps,
             Weight(product(idx(ps,i)^idx(t,i),i=0..size(ps)-1),
               Ret(ps))),
           dirichlet(ary(size(as),i,idx(as,i)+idx(t,i))),
           label="Conjugacy between rolled Dirichlet and multinomial (currently fails)");

kernelopts(assertlevel=2): # be strict on all assertions while testing
kernelopts(opaquemodules=false): # allow testing of internal routines
if not (NewSLO :: `module`) then
  WARNING("loading NewSLO failed");
  `quit`(3);
end if;

with(Hakaru):
with(NewSLO):

#####################################################################
#
# plate/array tests
#
#####################################################################

triv := Plate(k, i, Ret(i)):
TestHakaru(triv, Ret(ary(k, i, i)), label="Dirac Plate");

ary0 := Bind(Plate(k, i, Gaussian(0,1)), xs, Ret([xs])):
TestHakaru(ary0, ary0, label="plate roundtripping", ctx = KB:-assert(k::nonnegint,KB:-empty));

ary1  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(Unit))), ys,
         Ret(x))):
ary1w := 2^(-(1/2)*n+1/2)*exp((1/2)*((sum(idx(t,i),i=0..n-1))^2-(sum(idx(t,i)^2,i=0..n-1))*n-(sum(idx(t,i)^2,i=0..n-1)))/(n+1))*Pi^(-(1/2)*n)/sqrt(2+2*n):
TestHakaru(ary1,
  Weight(ary1w, Gaussian((sum(idx(t, i), i = 0 .. n-1))/(n+1), 1/sqrt(n+1))),
  label="Wednesday goal", ctx = KB:-assert(n::nonnegint,KB:-empty));
TestHakaru(Bind(ary1, x, Ret(Unit)), Weight(ary1w, Ret(Unit)),
  label="Wednesday goal total", ctx = KB:-assert(n::nonnegint,KB:-empty));
ary2  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Bind(Gaussian(idx(t,i),1),z, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(z+1)))), ys,
         Ret(ys))):
TestHakaru(ary2,
  Weight(ary1w, Bind(Plate(n, i, Gaussian(idx(t,i),1)), zs, Ret(ary(n, i, idx(zs,i)+1)))),
  label="Reason for fission", ctx = KB:-assert(n::nonnegint,KB:-empty));
ary3  := Bind(Gaussian(0,1), x,
         Bind(Plate(n, i, Bind(Gaussian(idx(t,i),1),z, Weight(density[Gaussian](x,1)(idx(t,i)), Ret(z)))), zs,
         Ret(zs))):
TestHakaru(ary3, Weight(ary1w, Plate(n, i, Gaussian(idx(t,i),1))),
  label="Array eta", ctx = KB:-assert(n::nonnegint,KB:-empty));
TestHakaru(Ret(ary(n,i,idx(f(i),i))), label="Don't overdo array eta");

bry1  := Bind(BetaD(alpha,beta), x,
         Bind(Plate(n, i, Weight(x    ^piecewise(idx(y,i)=true ,1) *
                                 (1-x)^piecewise(idx(y,i)=false,1),
                          Ret(Unit))), ys,
         Ret(x))):
bry1s := Weight(Beta(alpha+sum(piecewise(idx(y,i)=true ,1), i=0..n-1),
                     beta +sum(piecewise(idx(y,i)=false,1), i=0..n-1))/Beta(alpha,beta),
         BetaD(alpha+sum(piecewise(idx(y,i)=true ,1), i=0..n-1),
               beta +sum(piecewise(idx(y,i)=false,1), i=0..n-1))):
TestHakaru(bry1, bry1s,
  label="first way to express flipping a biased coin many times",
  ctx = KB:-assert(n::nonnegint,KB:-empty));

bry2  := Bind(BetaD(alpha,beta), x,
         Bind(Plate(n, i, Weight(x    ^(  idx(y,i)) *
                                 (1-x)^(1-idx(y,i)),
                          Ret(Unit))), ys,
         Ret(x))):
bry2s := Weight(Beta(alpha+  sum(idx(y,i),i=0..n-1),
                     beta +n-sum(idx(y,i),i=0..n-1))/Beta(alpha,beta),
         BetaD(alpha+  sum(idx(y,i),i=0..n-1),
               beta +n-sum(idx(y,i),i=0..n-1))):
TestHakaru(bry2, bry2s,
  label="second way to express flipping a biased coin many times",
  ctx = KB:-assert(n::nonnegint,KB:-empty));

fission     := Bind(Plate(k, i, Gaussian(z^i,1)), xs, Plate(k, i, Gaussian(idx(xs,i),1))):
fusion      := Plate(k, i, Bind(Gaussian(z^i,1), x, Gaussian(x,1))):
conjugacies := Plate(k, i, Gaussian(z^i, sqrt(2))):
conjugacies5:= Bind(Gaussian(z^0, sqrt(2)), x0,
               Bind(Gaussian(z^1, sqrt(2)), x1,
               Bind(Gaussian(z^2, sqrt(2)), x2,
               Bind(Gaussian(z^3, sqrt(2)), x3,
               Bind(Gaussian(z^4, sqrt(2)), x4,
               Ret(ary(5, i, piecewise(i=1,x1, i=2,x2, i=3,x3, i=4,x4, x0)))))))):
TestHakaru(eval(fission,{    z=1}), eval(conjugacies ,z=1), verify=normal, label="Conjugacy across iid plates");
TestHakaru(eval(fusion ,{    z=1}), eval(conjugacies ,z=1), verify=normal, label="Conjugacy in iid plate");
TestHakaru(eval(fission,{k=5,z=1}), eval(conjugacies5,z=1), verify=normal, label="Conjugacy across iid plates unrolled");
TestHakaru(eval(fusion ,{k=5,z=1}), eval(conjugacies5,z=1), verify=normal, label="Conjugacy in iid plate unrolled");
TestHakaru(     fission           ,      conjugacies      , verify=normal, label="Conjugacy across plates", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(     fusion            ,      conjugacies      , verify=normal, label="Conjugacy in plate", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(eval(fission,{k=5    }),      conjugacies5     , verify=normal, label="Conjugacy across plates unrolled", ctx=KB:-assert(z>0,KB:-empty));
TestHakaru(eval(fusion ,{k=5    }),      conjugacies5     , verify=normal, label="Conjugacy in plate unrolled", ctx=KB:-assert(z>0,KB:-empty));

# Simplifying gmm below is a baby step towards index manipulations we need
gmm := Bind(Plate(k, c, Gaussian(0,1)), xs,
       Bind(Plate(n, i, Weight(density[Gaussian](idx(xs,idx(cs,i)),1)(idx(t,i)), Ret(Unit))), ys,
       Ret(xs))):
gmm_s := Weight(2^(-(1/2)*n)*Pi^(-(1/2)*n)*exp(-(1/2)*(sum(idx(t,i)^2, i=0..n-1))),
         Bind(Plate(k, c, Gaussian(0, 1)), xs,
         Weight(exp(sum(idx(t,i)*idx(xs,idx(cs,i)), i=0..n-1))*exp(-(1/2)*(sum(idx(xs,idx(cs,i))^2, i=0..n-1))),
         Ret(xs)))):
TestHakaru(gmm, gmm_s, label="gmm");

# Detecting Dirichlet-multinomial conjugacy when unrolled
dirichlet := proc(as)
  Bind(Plate(size(as)-1, i, BetaD(sum(idx(as,j), j=0..i), idx(as,i+1))), xs,
  Ret(ary(size(as), i,
          product(idx(xs,j), j=i..size(as)-2)
          * piecewise(0=i, 1, 1-idx(xs,i-1)))))
end proc:
categorical := proc(ps)
  Bind(Counting(0, size(ps)-1), i, Weight(idx(ps,i), Ret(i)))
end proc:
# The next line eta-expands CodeTools[Test] so that its arguments get evaluated
(proc() CodeTools[Test](_passed) end proc)
  (op(2, [unweight(
     fromLO(improve(toLO(
       Plate(n,k,
         Bind(dirichlet(ary(5,i,1)),ps,
           Weight(product(idx(ps,i)^idx(idx(t,k),i),i=0..size(ps)-1),
             Ret(ps))))))))]),
   fromLO(toLO(
     Plate(n,k,
       dirichlet(ary(5,i,1+idx(idx(t,k),i)))))),
   measure(simplify),
   label="Conjugacy between unrolled symmetric Dirichlet and multinomial");
# We'd like the test above to pass even if the count 5 becomes symbolic.
# Below is some progress towards this goal:
TestHakaru(dirichlet(as), label="Dirichlet symbolic prior roundtrip");
TestHakaru(Bind(dirichlet(as),ps,
             Weight(product(idx(ps,i)^idx(t,i),i=0..size(ps)-1),
               Ret(ps))),
           dirichlet(ary(size(as),i,idx(as,i)+idx(t,i))),
           label="Conjugacy between rolled Dirichlet and multinomial (currently fails)");


simplification of control piecewise in comments (does not work because of alpha for Int).  Use Not rather than not in assuming.  Also for assuming, do not bother with And, concatenation will do.  Simplify weight when both branches the same.

# we create 3 new binding forms.  Teach Maple (through depends and eval)
# about them.
# Both LO and Integrand bind from 1st arg to 2nd arg, whereas Bind
# binds from 2nd arg to 3rd arg.
`depends/LO` := proc(h, e, x) depends(e, x minus {h}) end proc:

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

# note that v _can_ occur in m1.
`depends/Bind` := proc(m1, v, m2, x)
  depends(m1, x) or depends(m2, x minus {v})
end proc:

generic_evalat := proc(vv, mm, eqs)
  local v, m, eqsRemain, subsEq, eq, vRename, funs;
  v, m := vv, mm;
  funs := map2(op, 0, indets(mm, 'function'));
  eqsRemain := remove((eq -> op(1,eq) = op(2,eq)), eqs);
  eqsRemain, subsEq := selectremove((eq -> type(op(1,eq), 'name')), eqsRemain);
  eqsRemain := select((eq -> op(1,eq) <> v and
    (depends(mm, op(1,eq)) or member(op(1,eq), funs))), eqsRemain);
  if depends(eqsRemain, v) then
    vRename := gensym(v);
    m := subs(v=vRename, m);
    v := vRename;
  end if;
  m := subs(subsEq,m);
  if nops(eqsRemain) > 0 then
    m := eval(m,eqsRemain);
  end if;
  v, m;
end proc:

`eval/Integrate` := proc(e, eqs)
  local v, m2;
  v, m2 := op(e);
  eval(op(0,e), eqs)(generic_evalat(v, m2, eqs))
end proc:

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

`eval/Bind` := proc(e, eqs)
  local m1, v, m2;
  m1, v, m2 := op(e);
  eval(op(0,e), eqs)(eval(m1, eqs), generic_evalat(v, m2, eqs))
end proc:

`eval/lam` := proc(e, eqs)
  local v, m2;
  v, m2 := op(e);
  eval(op(0,e), eqs)(generic_evalat(v, m2, eqs))
end proc:

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

# make gensym global, so that it can be shared with other 'global' routines
gensym := module()
  export ModuleApply;
  local gs_counter;
  gs_counter := 0;
  ModuleApply := proc(x::name)
    gs_counter := gs_counter + 1;
    x || gs_counter;
  end proc;
end module: # gensym

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

NewSLO := module ()
  option package;
  local t_pw, unweight, factorize,
        recognize, get_de, recognize_de, Diffop, Recognized,
        step2, myint, myint_pw, simp_pw, simp_Int, get_indicators,
        indicator, extract_dom,
        verify_measure;
  export Integrand, applyintegrand, app, lam,
         Lebesgue, Uniform, Gaussian, Cauchy, BetaD, GammaD, StudentT,
         Ret, Bind, Msum, Weight, LO, Indicator,
         HakaruToLO, integrate, LOToHakaru, unintegrate,
         TestHakaru, measure, density,
         Simplify;

  t_pw := 'specfunc(piecewise)';

# 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).

# TODO evalapply/Integrand instead of applyintegrand?
# TODO evalapply/{Ret,Bind,...} instead of integrate?!

  applyintegrand := proc(h, x)
    if h :: 'Integrand(name, anything)' then
      eval(op(2,h), op(1,h) = x)
    elif h :: procedure then
      h(x)
    else
      'procname(_passed)'
    end if
  end proc;

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

  HakaruToLO := proc(m)
    local h;
    h := gensym('h');
    LO(h, integrate(m, h))
  end proc;

  integrate := proc(m, h)
    local x, n, i;
    if m :: 'Lebesgue()' then
      x := gensym('xl');
      Int(applyintegrand(h, x),
          x=-infinity..infinity)
    elif m :: 'Uniform(anything, anything)' then
      x := gensym('xu');
      Int(applyintegrand(h, x) / (op(2,m)-op(1,m)),
          x=op(1,m)..op(2,m))
    elif m :: '{Gaussian(anything, anything),
                Cauchy  (anything, anything),
                StudentT(anything, anything, anything)}' then
      x := gensym('xg');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=-infinity..infinity)
    elif m :: 'BetaD(anything, anything)' then
      x := gensym('xb');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=0..1)
    elif m :: 'GammaD(anything, anything)' then
      x := gensym('xr');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=0..infinity)
    elif m :: 'Ret(anything)' then
      applyintegrand(h, op(1,m))
    elif m :: 'Bind(anything, name, anything)' then
      integrate(op(1,m), z -> integrate(eval(op(3,m), op(2,m) = z), h))
    elif m :: 'specfunc(Msum)' then
      `+`(op(map(integrate, [op(m)], h)))
    elif m :: 'Weight(anything, anything)' then
      op(1,m) * integrate(op(2,m), h)
    elif m :: t_pw then
      n := nops(m);
      piecewise(seq(`if`(i::even or i=n, integrate(op(i,m), h), op(i,m)),
                    i=1..n))
    elif m :: 'LO(name, anything)' then
      eval(op(2,m), op(1,m) = h)
    elif h :: procedure then
      x := gensym('xa');
      'integrate'(m, Integrand(x, h(x)))
    else
      'procname(_passed)'
    end if
  end proc;

# Step 2 of 3: computer algebra

  Simplify := proc(lo :: LO(name, anything))
    LO(op(1,lo), step2(op(2,lo), op(1,lo), []))
  end proc;

  # Walk through integrals and simplify, recursing through grammar
  # h - name of the linear operator above us
  # constraints - domain information
  step2 := proc(e, h, constraints)
    local ee, subintegral, w;

    if e :: Int(anything, name = range) then
      simp_Int(e, bound, constraints)
    elif e :: `*` then
      (subintegral, w) := selectremove(depends, e, h);
      maptype(`*`, simp_pw, w) * step2(subintegral, h, constraints)
    # control piecewise; doesn't work...
    # elif e :: t_pw then
      # if both branches are the same, don't bother branching
    #   if nops(e)=3 and verify_measure(op(2,e), op(3,e), 'equal') then
    #     op(2,e)
    #   else
    #     e
    #   end if
    else
      e
    end if;
  end proc;

  # if a weight term is piecewise, then
  # 1. check if it is in  fact an indicator function, and if so, convert
  # 2. check if all its branches are equal, if so simplify
  simp_pw := proc(e)
    if e::t_pw then
     if nops(e) = 3 then
        # indicator
        if op(3,e)=0 then indicator(op(1,e)) * op(2,e)
        elif op(2,e)=0 then indicator(Not(op(1,e))) * op(3,e)

        # same
        elif Testzero(op(2,e) - op(3,e)) then
          op(2,e)
        else e
        end if
      else
        e
      end if
    else
      e
    end if
  end proc;

  simp_Int := proc(e, bound, constraints)
    local ee, var, rng, dom_spec, rest, new_rng, left;

    var := op([2,1],e);
    rng := op([2,2],e);

    # first step through the integrand
    ee := step2(op(1,e), h, [var = rng, op(constraints)]);

    # then if there are domain restrictions, try to apply them
    (dom_spec, rest) := get_indicators(ee);

    if not (dom_spec = 1) then
      if depends(dom_spec, var) then
        # TODO: new_rng might be a union of ranges
        (new_rng, left) := extract_dom(dom_spec, var);
        if new_rng :: 'range' then
          rng := max(op(1,rng), op(1,new_rng)) ..
                 min(op(2,rng), op(2,new_rng));
          ee := left * rest;
        end if;
      end if;
    end if;

    myint(ee, var = rng, constraints)
  end proc;

  get_indicators := proc(e)
    local inds, rest;
    if e::`*` then
      (inds, rest) := selectremove(type, e, 'specfunc(Indicator)');
      if inds::`*` then error "need to merge indicators" end if;
      (inds, rest)
    elif e::'specfunc(Indicator)' then
      e, 1
    else
      1, e
    end if
  end proc;

  extract_dom := proc(spec, v)
    local s, lo, hi, i, rest;
    if spec::'specfunc(Indicator)' then
      s, rest := selectremove(depends, spec, v);
      if nops(s) <> 1 then error "expected a single range" end if;
      s := op(1,s);
      lo := -infinity; hi := infinity;
      for i in s do
        if i :: 'numeric < identical(v)' or i :: 'numeric <= identical(v)' then
          lo := max(lo, op(1,i));
        elif i :: 'identical(v) < numeric' or
             i :: 'identical(v) <= numeric' then
          hi := min(hi, op(2,i));
        elif depends(i, v) then
          error "encountered (%1) in domain spec", i;
        # else -- no need, as these are not relevant
        end if;
      end do;
      if rest = 'Indicator()' then rest := 1; end if;
      lo .. hi, rest;
    else
      FAIL, spec;
    end if;
  end proc;

  indicator := proc(b)
    local to_set;

    to_set := proc(a)
      if a :: 'specfunc'(And) then
        map(op @ to_set, {op(a)})
      elif a :: 'specop(anything, `and`)' then
        map(op @ to_set, {op(a)})
      elif a :: {`<`,`<=`} then
        {a}
      else
        error "unknown set specification (%1)", a
      end if;
    end proc;

    Indicator(to_set(b))
  end proc;

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

  Bind := proc(m, x, n)
    if n = 'Ret'(x) then
      m # monad law: right identity
    elif m :: 'Ret(anything)' then
      eval(n, x = op(1,m)) # monad law: left identity
    else
      'procname(_passed)'
    end if;
  end proc;

  Weight := proc(p, m)
    if p = 1 then
      m
    elif p = 0 then
      Msum()
    elif m :: 'Weight(anything, anything)' then
      Weight(p * op(1,m), op(2,m))
    else
      'procname(_passed)'
    end if;
  end proc;

  LOToHakaru := proc(lo :: LO(name, anything))
    local h;
    h := gensym(op(1,lo));
    unintegrate(h, eval(op(2,lo), op(1,lo) = h), [])
  end proc;

  unintegrate := proc(h :: name, integral, context :: list)
    local x, lo, hi, m, w, w0, recognition, subintegral,
          n, i, next_context, update_context;
    if integral :: 'And'('specfunc({Int,int})',
                         'anyfunc'('anything','name'='range'('freeof'(h)))) then
      x := gensym(op([2,1],integral));
      (lo, hi) := op(op([2,2],integral));
      next_context := [op(context), lo<x, x<hi];
      # TODO: enrich context with x (measure class lebesgue)
      subintegral := eval(op(1,integral), op([2,1],integral) = x);
      (w, m) := unweight(unintegrate(h, subintegral, next_context));
      recognition := recognize(w, x, lo, hi) assuming op(next_context);
      if recognition :: 'Recognized(anything, anything)' then
        # Recognition succeeded
        (w, w0) := factorize(op(2,recognition), x);
        Weight(w0, Bind(op(1,recognition), x, Weight(w, m)))
      else
        # Recognition failed
        (w, w0) := factorize(w, x);
        m := Weight(w, m);
        if hi <> infinity then
          m := piecewise(x < hi, m, Msum())
        end if;
        if lo <> -infinity then
          m := piecewise(lo < x, m, Msum())
        end if;
        Weight(w0, Bind(Lebesgue(), x, m))
      end if
    elif integral :: 'applyintegrand'('identical'(h), 'freeof'(h)) then
      Ret(op(2,integral))
    elif integral = 0 then
      Msum()
    elif integral :: `+` then
      Msum(op(map2(unintegrate, h, convert(integral, 'list'), context)))
    elif integral :: `*` then
      (subintegral, w) := selectremove(depends, integral, h);
      if subintegral :: `*` then
        error "Nonlinear integral %1", integral
      end if;
      Weight(w, unintegrate(h, subintegral, context))
    elif integral :: t_pw
         and `and`(seq(not (depends(op(i,integral), h)),
                       i=1..nops(integral)-1, 2)) then
      n := nops(integral);
      next_context := context;
      update_context := proc(c)
        local then_context;
        then_context := [op(next_context), c];
        next_context := [op(next_context), Not(c)]; # Mutation!
        then_context
      end proc;
      piecewise(seq(piecewise(i::even,
                              unintegrate(h, op(i,integral),
                                          update_context(op(i-1,integral))),
                              i=n,
                              unintegrate(h, op(i,integral), next_context),
                              op(i,integral)),
                    i=1..n))
    elif integral :: 'integrate'('freeof'(h), 'anything') then
      x := gensym('x');
      # TODO is there any way to enrich context in this case?
      Bind(op(1,integral), x,
           unintegrate(h, applyintegrand(op(2,integral), x), context))
    else
      # Failure: return residual LO
      LO(h, integral)
    end if
  end proc;

  app := proc (func, argu)
    if func :: lam(name, anything) then
      eval(op(2,func), op(1,func)=argu)
    else
      'procname(_passed)'
    end if
  end proc;

  unweight := proc(m)
    local unweights, total;
    if m :: 'Weight(anything, anything)' then
      op(m)
    elif m :: 'specfunc(Msum)' then
      total := `+`(op(map((mi -> unweight(mi)[1]), m)));
      (total, map((mi -> Weight(1/total, mi)), m))
    else
      # TODO: Better weight estimate for piecewise & density-recognition cases?
      (1, m)
    end if;
  end proc;

  factorize := proc(weight, x)
    # return (weight, 1); # uncomment this to disable factorization
    if weight :: `*` then
      selectremove(depends, weight, x)
    elif depends(weight, x) then
      (weight, 1)
    else
      (1, weight)
    end if
  end proc;

  recognize := proc(weight, x, lo, hi)
    local de, Dx, f, w, res, rng;
    res := FAIL;
    de := get_de(weight, x, Dx, f);
    if de :: 'Diffop(anything, anything)' then
      res := recognize_de(op(de), Dx, f, x, lo, hi)
    end if;
    if res = FAIL then
      rng := hi - lo;
      w := simplify(weight * (hi - lo));
      # weight could be piecewise and simplify will hide the problem
      if not (rng :: 'SymbolicInfinity' or w :: {'SymbolicInfinity', 'undefined'}) then
        res := Recognized(Uniform(lo, hi), w)
      end if
    end if;
    res
  end proc;

  get_de := proc(dens, var, Dx, f)
    :: Or(Diffop(anything, set(function=anything)), identical(FAIL));
    local de, init;
    try
      de := gfun[holexprtodiffeq](dens, f(var));
      de := gfun[diffeqtohomdiffeq](de, f(var));
      if not (de :: set) then
        de := {de}
      end if;
      init, de := selectremove(type, de, `=`);
      if nops(de) = 1 then
        if nops(init) = 0 then
          # TODO: Replace {0, 1/2, 1} by PyMC's distribution-specific "testval"
          init := map((val -> f(val) = eval(dens, var=val)), {0, 1/2, 1})
        end if;
        return Diffop(DEtools[de2diffop](de[1], f(var), [Dx, var]), init)
      end if
    catch: # do nothing
    end try;
    FAIL
  end proc;

  recognize_de := proc(diffop, init, Dx, f, var, lo, hi)
    local dist, ii, constraints, w, a0, a1, a, b0, b1, c0, c1, c2, loc, nu;
    dist := FAIL;
    if lo = -infinity and hi = infinity
       and ispoly(diffop, 'linear', Dx, 'a0', 'a1') then
      a := normal(a0/a1);
      if ispoly(a, 'linear', var, 'b0', 'b1') then
        dist := Gaussian(-b0/b1, sqrt(1/b1))
      elif ispoly(numer(a), 'linear', var, 'b0', 'b1') and
           ispoly(denom(a), 'quadratic', var, 'c0', 'c1', 'c2') then
        loc := -c1/c2/2;
        nu := b1/c2 - 1;
        if Testzero(b0 + loc * b1) then
          if Testzero(nu - 1) then
            dist := Cauchy(loc, sqrt(c0/c2-loc^2))
          else
            dist := StudentT(nu, loc, sqrt((c0/c2-loc^2)/nu))
          end if
        end if
      end if;
    elif lo = 0 and hi = 1
         and ispoly(diffop, 'linear', Dx, 'a0', 'a1')
         and ispoly(normal(a0*var*(1-var)/a1), 'linear', var, 'b0', 'b1') then
      dist := BetaD(1-b0, 1+b0+b1)
    elif lo = 0 and hi = infinity
         and ispoly(diffop, 'linear', Dx, 'a0', 'a1')
         and ispoly(normal(a0*var/a1), 'linear', var, 'b0', 'b1') then
      dist := GammaD(1-b0, 1/b1)
    end if;
    if dist <> FAIL then
      try
        ii := map(convert, init, 'diff');
        constraints := eval(ii, f = (x -> w*density[op(0,dist)](op(dist))(x)));
        w := eval(w, solve(simplify(constraints), w));
        if not (has(w, 'w')) then
          return Recognized(dist, w)
        end if
      catch: # do nothing
      end try;
    end if;
    FAIL
  end proc;

  density[Gaussian] := proc(mu, sigma) proc(x)
    1/sigma/sqrt(2)/sqrt(Pi)*exp(-(x-mu)^2/2/sigma^2)
  end proc end proc;
  density[Cauchy] := proc(loc, scale) proc(x)
    1/Pi/scale/(1+((x-loc)/scale)^2)
  end proc end proc;
  density[StudentT] := proc(nu, mu, sigma) proc(x)
    GAMMA((nu+1)/2) / GAMMA(nu/2) / sqrt(Pi*nu) / sigma
    * (1 + ((x-mu)/sigma)^2/nu)^(-(nu+1)/2)
  end proc end proc;
  density[BetaD] := proc(a, b) proc(x)
    x^(a-1)*(1-x)^(b-1)/Beta(a,b)
  end proc end proc;
  # Hakaru uses the alternate definition of gamma, so the args are backwards
  density[GammaD] := proc(shape, scale) proc(x)
    x^(shape-1)/scale^shape*exp(-x/scale)/GAMMA(shape);
  end proc end proc;

  # a special version of int to get around Maple weaknesses
  myint := proc(expr, b, constraints)
    local var, inds, res;
    var := op(1,b);

    # grab all the pieces under arbitrary functions
    inds := map2(op, 2, indets(expr, 'specfunc(anything, applyintegrand)'));
    # trivial integration bounds
    if Normalizer(op([2,1],b)-op([2,2],b)) = 0 then
      res := 0;
    # if the integration variable doesn't occur under any of them:
    elif not depends(inds, var) then
      # just do it!
      res := int(expr, b)
    elif type(expr, 'piecewise') then
      res := myint_pw(expr, b); # TODO pass constraints too?
    else # give up
      res := Int(expr, b)
    end if;
    res;
  end proc;

  # special code for integrating pw
  myint_pw := proc(expr, b :: name = `..`)
    local rels, n, rest, i, res, res2, var, lower, upper, cond;
    n := floor(nops(expr)/2);
    rels := [seq(op(2*i-1, expr), i=1..n)];
    rest := evalb(2*n < nops(expr));
    upper := op([2,2],b);
    if type(rels, list({`<`, `<=`})) and indets(rels,'name') = {op(1,b)} then
      res := 0; lower := op([2,1], b); var := op(1,b);
      for i from 1 to n do
        cond := op(2*i-1, expr);
        if cond::{identical(var) < anything, identical(var) <= anything} and
          evalb(signum(op(2,cond) - lower) = 1) then # op2 > lower
          res := res + myint(op(2*i, expr), var = lower .. min(upper,op(2,cond)));
          lower := op(2,cond);
        elif cond::{anything < identical(var), anything <= identical(var)} and
          evalb(signum(upper - op(1,cond)) = 1) then # op1 < upper
          res := res + myint(op(2*i, expr), var = max(lower,op(1,cond)) .. upper);
          upper := op(1,cond);
        else
          error "cannot handle condition (%1) while integrating pw", cond;
        end if;
      end do;
      if rest then # note that lower=upper is possible, but ok
        if lower < upper then
          res := res + myint(op(-1, expr), var = lower..upper);
        end if;
      end if;
      res
    elif length(expr) < 200 then # what have we got to lose?
      int(expr, b);
    else # well, too much time!
      Int(expr, b)
    end if;
  end proc;

# Testing

  TestHakaru := proc(m,n:=m,{simp:=Simplify})
    CodeTools[Test](LOToHakaru(simp(HakaruToLO(m))), n,
      measure({boolean,equal}), _rest)
  end proc;

  verify_measure := proc(m, n, v:='boolean')
    local mv, x, i, j, k;
    mv := measure(v);
    if verify(m, n, 'Bind'(mv, true, true)) then
      x := gensym(cat(op(2,m), "_", op(2,n), "_"));
      thisproc(subs(op(2,m)=x, op(3,m)),
               subs(op(2,n)=x, op(3,n)), v)
    elif m :: 'specfunc(Msum)' and n :: 'specfunc(Msum)'
         and nops(m) = nops(n) then
      k := nops(m);
      verify(k, GraphTheory[BipartiteMatching](GraphTheory[Graph]({
                seq(seq(`if`(thisproc(op(i,m), op(j,n), v), {i,-j}, NULL),
                        j=1..k), i=1..k)}))[1]);
    elif m :: t_pw and n :: t_pw and nops(m) = nops(n) then
      k := nops(m);
      for i to k do
        if not (thisproc(op(i,m), op(i,n), `if`(i::even or i=k, mv, v))) then
          return false
        end if
      end do;
      true
    elif m :: 'LO(name, anything)' and n :: 'LO(name, anything)' then
      x := gensym(cat(op(1,m), "_", op(1,n), "_"));
      verify(subs(op(1,m)=x, op(2,m)),
             subs(op(1,n)=x, op(2,n)), v)
    else
      verify(m, n, {v,
        Lebesgue(),
        Uniform(v, v),
        Gaussian(v, v),
        Cauchy(v, v),
        StudentT(v, v, v),
        BetaD(v, v),
        GammaD(v, v),
        Ret(mv),
        Weight(v, mv)
      })
    end if
  end proc;

  VerifyTools[AddVerification](measure = verify_measure);

end module; # NewSLO

# we create 3 new binding forms.  Teach Maple (through depends and eval)
# about them.
# Both LO and Integrand bind from 1st arg to 2nd arg, whereas Bind
# binds from 2nd arg to 3rd arg.
`depends/LO` := proc(h, e, x) depends(e, x minus {h}) end proc:

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

# note that v _can_ occur in m1.
`depends/Bind` := proc(m1, v, m2, x)
  depends(m1, x) or depends(m2, x minus {v})
end proc:

generic_evalat := proc(vv, mm, eqs)
  local v, m, eqsRemain, subsEq, eq, vRename, funs;
  v, m := vv, mm;
  funs := map2(op, 0, indets(mm, 'function'));
  eqsRemain := remove((eq -> op(1,eq) = op(2,eq)), eqs);
  eqsRemain, subsEq := selectremove((eq -> type(op(1,eq), 'name')), eqsRemain);
  eqsRemain := select((eq -> op(1,eq) <> v and
    (depends(mm, op(1,eq)) or member(op(1,eq), funs))), eqsRemain);
  if depends(eqsRemain, v) then
    vRename := gensym(v);
    m := subs(v=vRename, m);
    v := vRename;
  end if;
  m := subs(subsEq,m);
  if nops(eqsRemain) > 0 then
    m := eval(m,eqsRemain);
  end if;
  v, m;
end proc:

`eval/Integrate` := proc(e, eqs)
  local v, m2;
  v, m2 := op(e);
  eval(op(0,e), eqs)(generic_evalat(v, m2, eqs))
end proc:

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

`eval/Bind` := proc(e, eqs)
  local m1, v, m2;
  m1, v, m2 := op(e);
  eval(op(0,e), eqs)(eval(m1, eqs), generic_evalat(v, m2, eqs))
end proc:

`eval/lam` := proc(e, eqs)
  local v, m2;
  v, m2 := op(e);
  eval(op(0,e), eqs)(generic_evalat(v, m2, eqs))
end proc:

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

# make gensym global, so that it can be shared with other 'global' routines
gensym := module()
  export ModuleApply;
  local gs_counter;
  gs_counter := 0;
  ModuleApply := proc(x::name)
    gs_counter := gs_counter + 1;
    x || gs_counter;
  end proc;
end module: # gensym

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

NewSLO := module ()
  option package;
  local t_pw, unweight, factorize,
        recognize, get_de, recognize_de, Diffop, Recognized,
        step2, myint, myint_pw, simp_pw, simp_Int, get_indicators,
        indicator, extract_dom,
        verify_measure;
  export Integrand, applyintegrand, app, lam,
         Lebesgue, Uniform, Gaussian, Cauchy, BetaD, GammaD, StudentT,
         Ret, Bind, Msum, Weight, LO, Indicator,
         HakaruToLO, integrate, LOToHakaru, unintegrate,
         TestHakaru, measure, density,
         Simplify;

  t_pw := 'specfunc(piecewise)';

# 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).

# TODO evalapply/Integrand instead of applyintegrand?
# TODO evalapply/{Ret,Bind,...} instead of integrate?!

  applyintegrand := proc(h, x)
    if h :: 'Integrand(name, anything)' then
      eval(op(2,h), op(1,h) = x)
    elif h :: procedure then
      h(x)
    else
      'procname(_passed)'
    end if
  end proc;

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

  HakaruToLO := proc(m)
    local h;
    h := gensym('h');
    LO(h, integrate(m, h))
  end proc;

  integrate := proc(m, h)
    local x, n, i;
    if m :: 'Lebesgue()' then
      x := gensym('xl');
      Int(applyintegrand(h, x),
          x=-infinity..infinity)
    elif m :: 'Uniform(anything, anything)' then
      x := gensym('xu');
      Int(applyintegrand(h, x) / (op(2,m)-op(1,m)),
          x=op(1,m)..op(2,m))
    elif m :: '{Gaussian(anything, anything),
                Cauchy  (anything, anything),
                StudentT(anything, anything, anything)}' then
      x := gensym('xg');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=-infinity..infinity)
    elif m :: 'BetaD(anything, anything)' then
      x := gensym('xb');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=0..1)
    elif m :: 'GammaD(anything, anything)' then
      x := gensym('xr');
      Int(density[op(0,m)](op(m))(x) * applyintegrand(h, x),
          x=0..infinity)
    elif m :: 'Ret(anything)' then
      applyintegrand(h, op(1,m))
    elif m :: 'Bind(anything, name, anything)' then
      integrate(op(1,m), z -> integrate(eval(op(3,m), op(2,m) = z), h))
    elif m :: 'specfunc(Msum)' then
      `+`(op(map(integrate, [op(m)], h)))
    elif m :: 'Weight(anything, anything)' then
      op(1,m) * integrate(op(2,m), h)
    elif m :: t_pw then
      n := nops(m);
      piecewise(seq(`if`(i::even or i=n, integrate(op(i,m), h), op(i,m)),
                    i=1..n))
    elif m :: 'LO(name, anything)' then
      eval(op(2,m), op(1,m) = h)
    elif h :: procedure then
      x := gensym('xa');
      'integrate'(m, Integrand(x, h(x)))
    else
      'procname(_passed)'
    end if
  end proc;

# Step 2 of 3: computer algebra

  Simplify := proc(lo :: LO(name, anything))
    LO(op(1,lo), step2(op(2,lo), op(1,lo), []))
  end proc;

  # Walk through integrals and simplify, recursing through grammar
  # h - name of the linear operator above us
  # constraints - domain information
  step2 := proc(e, h, constraints)
    local ee, subintegral, w;

    if e :: Int(anything, name = range) then
      simp_Int(e, bound, constraints)
    elif e :: `*` then
      (subintegral, w) := selectremove(depends, e, h);
      maptype(`*`, simp_pw, w) * subintegral
    else
      e
    end if;
  end proc;

  # if a weight term is piecewise, then check if it is in
  # fact an indicator function, and if so, convert
  simp_pw := proc(e)
    if e::t_pw then
     if nops(e) = 3 then
        if op(3,e)=0 then indicator(op(1,e)) * op(2,e)
        elif op(2,e)=0 then indicator(Not(op(1,e))) * op(3,e)
        else e
        end if
      else
        e
      end if
    else
      e
    end if
  end proc;

  simp_Int := proc(e, bound, constraints)
    local ee, var, rng, dom_spec, rest, new_rng, left;

    var := op([2,1],e);
    rng := op([2,2],e);

    # first step through the integrand
    ee := step2(op(1,e), h, [var = rng, op(constraints)]);

    # then if there are domain restrictions, try to apply them
    (dom_spec, rest) := get_indicators(ee);

    if not (dom_spec = 1) then
      if depends(dom_spec, var) then
        # TODO: new_rng might be a union of ranges
        (new_rng, left) := extract_dom(dom_spec, var);
        if new_rng :: 'range' then
          rng := max(op(1,rng), op(1,new_rng)) ..
                 min(op(2,rng), op(2,new_rng));
          ee := left * rest;
        end if;
      end if;
    end if;

    myint(ee, var = rng, constraints)
  end proc;

  get_indicators := proc(e)
    local inds, rest;
    if e::`*` then
      (inds, rest) := selectremove(type, e, 'specfunc(Indicator)');
      if inds::`*` then error "need to merge indicators" end if;
      (inds, rest)
    elif e::'specfunc(Indicator)' then
      e, 1
    else
      1, e
    end if
  end proc;

  extract_dom := proc(spec, v)
    local s, lo, hi, i, rest;
    if spec::'specfunc(Indicator)' then
      s, rest := selectremove(depends, spec, v);
      if nops(s) <> 1 then error "expected a single range" end if;
      s := op(1,s);
      lo := -infinity; hi := infinity;
      for i in s do
        if i :: 'numeric < identical(v)' or i :: 'numeric <= identical(v)' then
          lo := max(lo, op(1,i));
        elif i :: 'identical(v) < numeric' or
             i :: 'identical(v) <= numeric' then
          hi := min(hi, op(2,i));
        elif depends(i, v) then
          error "encountered (%1) in domain spec", i;
        # else -- no need, as these are not relevant
        end if;
      end do;
      if rest = 'Indicator()' then rest := 1; end if;
      lo .. hi, rest;
    else
      FAIL, spec;
    end if;
  end proc;

  indicator := proc(b)
    local to_set;

    to_set := proc(a)
      if a :: 'specfunc'(And) then
        map(op @ to_set, {op(a)})
      elif a :: 'specop(anything, `and`)' then
        map(op @ to_set, {op(a)})
      elif a :: {`<`,`<=`} then
        {a}
      else
        error "unknown set specification (%1)", a
      end if;
    end proc;

    Indicator(to_set(b))
  end proc;

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

  Bind := proc(m, x, n)
    if n = 'Ret'(x) then
      m # monad law: right identity
    elif m :: 'Ret(anything)' then
      eval(n, x = op(1,m)) # monad law: left identity
    else
      'procname(_passed)'
    end if;
  end proc;

  Weight := proc(p, m)
    if p = 1 then
      m
    elif p = 0 then
      Msum()
    elif m :: 'Weight(anything, anything)' then
      Weight(p * op(1,m), op(2,m))
    else
      'procname(_passed)'
    end if;
  end proc;

  LOToHakaru := proc(lo :: LO(name, anything))
    local h;
    h := gensym(op(1,lo));
    unintegrate(h, eval(op(2,lo), op(1,lo) = h), [])
  end proc;

  unintegrate := proc(h :: name, integral, context :: list)
    local x, lo, hi, m, w, w0, recognition, subintegral,
          n, i, next_context, update_context;
    if integral :: 'And'('specfunc({Int,int})',
                         'anyfunc'('anything','name'='range'('freeof'(h)))) then
      x := gensym(op([2,1],integral));
      (lo, hi) := op(op([2,2],integral));
      next_context := [op(context), And(lo<x, x<hi)];
      # TODO: enrich context with x (measure class lebesgue)
      subintegral := eval(op(1,integral), op([2,1],integral) = x);
      (w, m) := unweight(unintegrate(h, subintegral, next_context));
      recognition := recognize(w, x, lo, hi) assuming op(next_context);
      if recognition :: 'Recognized(anything, anything)' then
        # Recognition succeeded
        (w, w0) := factorize(op(2,recognition), x);
        Weight(w0, Bind(op(1,recognition), x, Weight(w, m)))
      else
        # Recognition failed
        (w, w0) := factorize(w, x);
        m := Weight(w, m);
        if hi <> infinity then
          m := piecewise(x < hi, m, Msum())
        end if;
        if lo <> -infinity then
          m := piecewise(lo < x, m, Msum())
        end if;
        Weight(w0, Bind(Lebesgue(), x, m))
      end if
    elif integral :: 'applyintegrand'('identical'(h), 'freeof'(h)) then
      Ret(op(2,integral))
    elif integral = 0 then
      Msum()
    elif integral :: `+` then
      Msum(op(map2(unintegrate, h, convert(integral, 'list'), context)))
    elif integral :: `*` then
      (subintegral, w) := selectremove(depends, integral, h);
      if subintegral :: `*` then
        error "Nonlinear integral %1", integral
      end if;
      Weight(w, unintegrate(h, subintegral, context))
    elif integral :: t_pw
         and `and`(seq(not (depends(op(i,integral), h)),
                       i=1..nops(integral)-1, 2)) then
      n := nops(integral);
      next_context := context;
      update_context := proc(c)
        local then_context;
        then_context := [op(next_context), c];
        next_context := [op(next_context), not c]; # Mutation!
        then_context
      end proc;
      piecewise(seq(piecewise(i::even,
                              unintegrate(h, op(i,integral),
                                          update_context(op(i-1,integral))),
                              i=n,
                              unintegrate(h, op(i,integral), next_context),
                              op(i,integral)),
                    i=1..n))
    elif integral :: 'integrate'('freeof'(h), 'anything') then
      x := gensym('x');
      # TODO is there any way to enrich context in this case?
      Bind(op(1,integral), x,
           unintegrate(h, applyintegrand(op(2,integral), x), context))
    else
      # Failure: return residual LO
      LO(h, integral)
    end if
  end proc;

  app := proc (func, argu)
    if func :: lam(name, anything) then
      eval(op(2,func), op(1,func)=argu)
    else
      'procname(_passed)'
    end if
  end proc;

  unweight := proc(m)
    local unweights, total;
    if m :: 'Weight(anything, anything)' then
      op(m)
    elif m :: 'specfunc(Msum)' then
      total := `+`(op(map((mi -> unweight(mi)[1]), m)));
      (total, map((mi -> Weight(1/total, mi)), m))
    else
      # TODO: Better weight estimate for piecewise & density-recognition cases?
      (1, m)
    end if;
  end proc;

  factorize := proc(weight, x)
    # return (weight, 1); # uncomment this to disable factorization
    if weight :: `*` then
      selectremove(depends, weight, x)
    elif depends(weight, x) then
      (weight, 1)
    else
      (1, weight)
    end if
  end proc;

  recognize := proc(weight, x, lo, hi)
    local de, Dx, f, w, res, rng;
    res := FAIL;
    de := get_de(weight, x, Dx, f);
    if de :: 'Diffop(anything, anything)' then
      res := recognize_de(op(de), Dx, f, x, lo, hi)
    end if;
    if res = FAIL then
      rng := hi - lo;
      w := simplify(weight * (hi - lo));
      # weight could be piecewise and simplify will hide the problem
      if not (rng :: 'SymbolicInfinity' or w :: {'SymbolicInfinity', 'undefined'}) then
        res := Recognized(Uniform(lo, hi), w)
      end if
    end if;
    res
  end proc;

  get_de := proc(dens, var, Dx, f)
    :: Or(Diffop(anything, set(function=anything)), identical(FAIL));
    local de, init;
    try
      de := gfun[holexprtodiffeq](dens, f(var));
      de := gfun[diffeqtohomdiffeq](de, f(var));
      if not (de :: set) then
        de := {de}
      end if;
      init, de := selectremove(type, de, `=`);
      if nops(de) = 1 then
        if nops(init) = 0 then
          # TODO: Replace {0, 1/2, 1} by PyMC's distribution-specific "testval"
          init := map((val -> f(val) = eval(dens, var=val)), {0, 1/2, 1})
        end if;
        return Diffop(DEtools[de2diffop](de[1], f(var), [Dx, var]), init)
      end if
    catch: # do nothing
    end try;
    FAIL
  end proc;

  recognize_de := proc(diffop, init, Dx, f, var, lo, hi)
    local dist, ii, constraints, w, a0, a1, a, b0, b1, c0, c1, c2, loc, nu;
    dist := FAIL;
    if lo = -infinity and hi = infinity
       and ispoly(diffop, 'linear', Dx, 'a0', 'a1') then
      a := normal(a0/a1);
      if ispoly(a, 'linear', var, 'b0', 'b1') then
        dist := Gaussian(-b0/b1, sqrt(1/b1))
      elif ispoly(numer(a), 'linear', var, 'b0', 'b1') and
           ispoly(denom(a), 'quadratic', var, 'c0', 'c1', 'c2') then
        loc := -c1/c2/2;
        nu := b1/c2 - 1;
        if Testzero(b0 + loc * b1) then
          if Testzero(nu - 1) then
            dist := Cauchy(loc, sqrt(c0/c2-loc^2))
          else
            dist := StudentT(nu, loc, sqrt((c0/c2-loc^2)/nu))
          end if
        end if
      end if;
    elif lo = 0 and hi = 1
         and ispoly(diffop, 'linear', Dx, 'a0', 'a1')
         and ispoly(normal(a0*var*(1-var)/a1), 'linear', var, 'b0', 'b1') then
      dist := BetaD(1-b0, 1+b0+b1)
    elif lo = 0 and hi = infinity
         and ispoly(diffop, 'linear', Dx, 'a0', 'a1')
         and ispoly(normal(a0*var/a1), 'linear', var, 'b0', 'b1') then
      dist := GammaD(1-b0, 1/b1)
    end if;
    if dist <> FAIL then
      try
        ii := map(convert, init, 'diff');
        constraints := eval(ii, f = (x -> w*density[op(0,dist)](op(dist))(x)));
        w := eval(w, solve(simplify(constraints), w));
        if not (has(w, 'w')) then
          return Recognized(dist, w)
        end if
      catch: # do nothing
      end try;
    end if;
    FAIL
  end proc;

  density[Gaussian] := proc(mu, sigma) proc(x)
    1/sigma/sqrt(2)/sqrt(Pi)*exp(-(x-mu)^2/2/sigma^2)
  end proc end proc;
  density[Cauchy] := proc(loc, scale) proc(x)
    1/Pi/scale/(1+((x-loc)/scale)^2)
  end proc end proc;
  density[StudentT] := proc(nu, mu, sigma) proc(x)
    GAMMA((nu+1)/2) / GAMMA(nu/2) / sqrt(Pi*nu) / sigma
    * (1 + ((x-mu)/sigma)^2/nu)^(-(nu+1)/2)
  end proc end proc;
  density[BetaD] := proc(a, b) proc(x)
    x^(a-1)*(1-x)^(b-1)/Beta(a,b)
  end proc end proc;
  # Hakaru uses the alternate definition of gamma, so the args are backwards
  density[GammaD] := proc(shape, scale) proc(x)
    x^(shape-1)/scale^shape*exp(-x/scale)/GAMMA(shape);
  end proc end proc;

  # a special version of int to get around Maple weaknesses
  myint := proc(expr, b, constraints)
    local var, inds, res;
    var := op(1,b);

    # grab all the pieces under arbitrary functions
    inds := map2(op, 2, indets(expr, 'specfunc(anything, applyintegrand)'));
    # trivial integration bounds
    if Normalizer(op([2,1],b)-op([2,2],b)) = 0 then
      res := 0;
    # if the integration variable doesn't occur under any of them:
    elif not depends(inds, var) then
      # just do it!
      res := int(expr, b)
    elif type(expr, 'piecewise') then
      res := myint_pw(expr, b); # TODO pass constraints too?
    else # give up
      res := Int(expr, b)
    end if;
    res;
  end proc;

  # special code for integrating pw
  myint_pw := proc(expr, b :: name = `..`)
    local rels, n, rest, i, res, res2, var, lower, upper, cond;
    n := floor(nops(expr)/2);
    rels := [seq(op(2*i-1, expr), i=1..n)];
    rest := evalb(2*n < nops(expr));
    upper := op([2,2],b);
    if type(rels, list({`<`, `<=`})) and indets(rels,'name') = {op(1,b)} then
      res := 0; lower := op([2,1], b); var := op(1,b);
      for i from 1 to n do
        cond := op(2*i-1, expr);
        if cond::{identical(var) < anything, identical(var) <= anything} and
          evalb(signum(op(2,cond) - lower) = 1) then # op2 > lower
          res := res + myint(op(2*i, expr), var = lower .. min(upper,op(2,cond)));
          lower := op(2,cond);
        elif cond::{anything < identical(var), anything <= identical(var)} and
          evalb(signum(upper - op(1,cond)) = 1) then # op1 < upper
          res := res + myint(op(2*i, expr), var = max(lower,op(1,cond)) .. upper);
          upper := op(1,cond);
        else
          error "cannot handle condition (%1) while integrating pw", cond;
        end if;
      end do;
      if rest then # note that lower=upper is possible, but ok
        if lower < upper then
          res := res + myint(op(-1, expr), var = lower..upper);
        end if;
      end if;
      res
    elif length(expr) < 200 then # what have we got to lose?
      int(expr, b);
    else # well, too much time!
      Int(expr, b)
    end if;
  end proc;

# Testing

  TestHakaru := proc(m,n:=m,{simp:=Simplify})
    CodeTools[Test](LOToHakaru(simp(HakaruToLO(m))), n,
      measure({boolean,equal}), _rest)
  end proc;

  verify_measure := proc(m, n, v:='boolean')
    local mv, x, i, j, k;
    mv := measure(v);
    if verify(m, n, 'Bind'(mv, true, true)) then
      x := gensym(cat(op(2,m), "_", op(2,n), "_"));
      thisproc(subs(op(2,m)=x, op(3,m)),
               subs(op(2,n)=x, op(3,n)), v)
    elif m :: 'specfunc(Msum)' and n :: 'specfunc(Msum)'
         and nops(m) = nops(n) then
      k := nops(m);
      verify(k, GraphTheory[BipartiteMatching](GraphTheory[Graph]({
                seq(seq(`if`(thisproc(op(i,m), op(j,n), v), {i,-j}, NULL),
                        j=1..k), i=1..k)}))[1]);
    elif m :: t_pw and n :: t_pw and nops(m) = nops(n) then
      k := nops(m);
      for i to k do
        if not (thisproc(op(i,m), op(i,n), `if`(i::even or i=k, mv, v))) then
          return false
        end if
      end do;
      true
    elif m :: 'LO(name, anything)' and n :: 'LO(name, anything)' then
      x := gensym(cat(op(1,m), "_", op(1,n), "_"));
      verify(subs(op(1,m)=x, op(2,m)),
             subs(op(1,n)=x, op(2,n)), v)
    else
      verify(m, n, {v,
        Lebesgue(),
        Uniform(v, v),
        Gaussian(v, v),
        Cauchy(v, v),
        StudentT(v, v, v),
        BetaD(v, v),
        GammaD(v, v),
        Ret(mv),
        Weight(v, mv)
      })
    end if
  end proc;

  VerifyTools[AddVerification](measure = verify_measure);

end module; # NewSLO