// particles-iso.diderot
//
// distributes particles on an isosurface
//

input string dataFile;           // name of dataset
input real radius = 0.2;         // radius of particle
input real isovalue = 0.0;       // isocontour value
input real cnstrEps = 0.01;      // epsilon test for constraint satisfaction
input real travelMax = 5.0;      // point can't move further than this
input int cnstrIterMax = 10;     // constraint satisfaction can't iterate more than this
input real stepLimit = 2;        // limit on how far a particle can move in one iteration
input real stepBack = 0.2;       // how to reduce step size if badstep
input real stepUp = 1.1;         // opportunistic scaling of step
input int pointNum = 100;        // number of initial seed points
input real energyConvg = 0.1;    // convergence threshold on system energy

image(3)[] img = load (dataFile);
field#2(3)[] F = img ⊛ bspln3;

// should allow functions to take actors, but sensible semantics
// limit that to a read-only view of actor at start of iteration;
// compiler can warn if the actor is passed post-killing-definition

// uses Newton-Raphson to find the zero-crossing
(bool, vec3) constrain (vec3 posOrig)
{
    vec3 grad, fix, pos = posOrig;
    int iter = 0;
    do {
        // implicitly a (iter,update,pos) tuple is iterated by this
        iter += 1;
        grad = ∇F@pos;
        fix = F@pos*grad / |grad|^2  // should optimize out sqrt()^2
        pos -= fix;
        if (dot(pos - posOrig,pos - posOrig) > travelMax*travelMax) {
           // ?can use "NaN" as a type real literal for IEEE Not-A-Number?
           return (false, [NaN,NaN,NaN])
        }
        if (iter > constrIterMax) {
           return (false, [NaN,NaN,NaN])
        }
    } while (dot(fix,fix) > cnstrEps*cnstrEps)
    // we've converged
    return (true,pos);
}

// this is the sort of function that would benefit from automatic 
// differentiation; it simply evaluates a piece-wise polynomial
// along with its derivative
(real, real) phi (real x)
{
    // this directly copied from C code (teem/src/pull/energy.c);
    real wx = 0.66, wy = 0.005;
    real enr, denr;  // output

    real xmo = wx-1;
    real xmoo = xmo*xmo;
    real xmooo = xmoo*xmo;
    if (x < wx) {
      real a = -3*(xmoo + (-1 + 2*wx)*wy)/(xmoo*wx);
      real b = 3*(xmoo + (-1 + wx*(2 + wx))*wy)/(xmoo*wx*wx);
      real c = (-1 + wy - wx*(-2 + wx + 2*(1 + wx)*wy))/(xmoo*wx*wx*wx);
      enr = 1 + x*(a + x*(b + x*c));
      denr = a + x*(2*b + x*3*c);
    } else if (x < 1) {
      real d = ((-1 + 3*wx)*wy)/xmooo;
      real e = -(6*wx*wy)/xmooo;
      real f = (3*(1 + wx)*wy)/xmooo;
      real g = -(2*wy)/xmooo;
      enr = d + x*(e + x*(f + x*g));
      denr = e + x*(2*f + x*3*g);
    } else {
      enr = 0;
      denr = 0;
    }
    return (enr, denr);
}

// evaluates the inter-particle energy function in 3D
(real, vec3) pairEnergy(vec3 here, vec3 there) {
    vec3 dir = (there - here)/|there - here|;
    (real enr, real denr) = phi(|there - here|/radius);
    return (enr, -denr*dir);
}

// computes energy and force from nearby particles
(real, vec3) totalEnergy(vec3 here) {
    // ?what is our sequence/list/iterable type?
    // ?can use ".pos" to refer to position variable?
    // ?neighbors has to know not return actor at "here" exactly?
    seq(vec3) neighPos = [n.pos for n in neighbors(radius, here)];
    real energy = 0;
    vec3 force = [0,0,0];
    for p in neighPos { 
        (energy, force) += pairEnergy(pos, n.pos); 
    }
    return (energy, force);
}

actor IsoPoint (vec3 posInit)
{
    real step = 1.0;   // initial step size for gradient descent
    output real energy;
    output vec3 pos;
    
    initially {
        (bool ok, vec3 pos) = constrain(posInit);
        if (!ok) {
            // we couldn't satisfy contraint from seed location; we're gone
            delete;   // ?what's the way for an actor to remove itself? 
        }
    }

    update
    {
        bool badstep;
        vec3 posNew, force;
        real energNew, energyOrig;

        (energyOrig, force) = totalEnergy(pos);
        // project force onto tangent plane
        force -= outer(∇F@pos, ∇F@pos)*force
        // ensures that step*force is a reasonable distance
        step = min(step, travelMax/|force|);
        if (step*|force| > travelMax) {
           step = travelMax/|force|;
        }
        do {
           (bool ok, posNew) = constrain(pos + step*force);
           if (!ok) {
               delete;
           }
           // there should be non-trivial savings from not computing force
           (energyNew,) = pointEnergy(pos, posNew);
           // use of badstep here is rather clunky
           badstep = energyNew > energyOrig;
           if (badstep) {
               step *= stepBack;
           }
        } while badstep;
        step *= stepUp;
        pos = posNew;
        energy = energyNew;
    }

    // (stabilization is not determined per-actor)
    
}

// ?from here on out I'm making things up ...?

// need some way of generating a list of random locations inside image,
// and then seeding actors there
initially [ IsoPoint(pos) for pos in randomPositionsInside(F, pointNum) ];

// keep update()ing all actors until total change in system is
// below some threshold
real Eold, Enew = reduce(sum, [a.energy for a in actors]);
do {
  Eold = Enew;
  updateAll;
  Enew = reduce(sum, [a.energy ]);
} while (Eold - Enew > energyConvg);