// quad-k1k2: demo of curvature calculation on quad patch dataset
// to process output:
// unu reshape -i mip.txt -s 3 200 200 | unu quantize -b 8 -o zslice-k1k2.png
//

int imgRes = 200;
int dataRes = 110;
vec3 orig = [0.0, 0.0, 10.0];
vec3 uVec = [1.0, 0.0, 0.0];
vec3 vVec = [0.0, 1.0, 0.0];
image(3)[] vol = load("../data/quad-patches-pad.nrrd");
field#2(3)[] F = vol ⊛ bspln3;
image(2)[3] imgRGB = load("../data/hue.nhdr");
field#0(2)[3] RGB = imgRGB ⊛ tent;

strand sample (int ui, int vi)
{
    real uu = lerp(-0.5, real(dataRes)-0.5, -0.5, real(ui), real(imgRes)-0.5);
    real vv = lerp(-0.5, real(dataRes)-0.5, -0.5, real(vi), real(imgRes)-0.5);
    output vec3 rgb = [0.0,0.0,0.0];

    update
    {
        vec3 pos = orig + uu*uVec + vv*vVec;
        if (inside (pos,F)) {
           vec3 g = -∇F(pos);
           // the "if (g[2] > 0.10)" test is specific to this dataset;
           // it is true when within one of the quadratic patches
           if (g[2] > 0.10) {
              vec3 n = normalize(g);
              tensor[3,3] H = -∇(∇F)(pos);
              tensor[3,3] P = identity[3] - n⊗n;
              tensor[3,3] G = (P•H•P)/|g|;
              real disc = max(0.0, sqrt(2.0*|G|^2 - trace(G)^2));
              real k1 = (trace(G) + disc)/2.0;
              real k2 = (trace(G) - disc)/2.0;
              // correct curvature computation ends here
              k1 = max(-1.0, min(1.0, 4.2*k1));
              k2 = max(-1.0, min(1.0, 4.2*k2));
              rgb = RGB([k1,k2]) if inside([k1,k2],RGB) else [0.0,0.0,0.0];
           }
        }
        stabilize;
    }
}

initially [ sample(ui, vi) | vi in 0..(imgRes-1), ui in 0..(imgRes-1) ];