/*! \file bmark-teem.c
 *
 * Gordon Kindlmann
 */

#if 1
#  define REAL  float
#  define ATAN2 atan2f
#  define SQRT  sqrtf
#else
#  define REAL  double
#  define ATAN2 atan2
#  define SQRT  sqrt
#endif

#define ELL_3M_FROBSQRD(m) \
  (ELL_3V_DOT((m)+0, (m)+0) + \
   ELL_3V_DOT((m)+3, (m)+3) + \
   ELL_3V_DOT((m)+6, (m)+6))


#include <stdio.h>
#include <math.h>
#include <stdbool.h>
#include "teem/nrrd.h"
#include "teem/gage.h"

#include "teem-defs.h"

STATIC_INLINE double dot3(double vec1[3], double vec2[3])
{
  return (vec1[0] * vec2[0] + vec1[1] * vec2[1] + vec1[2] * vec2[2]);
}

STATIC_INLINE double mag3(double vec[3])
{
  return sqrt(dot3(vec, vec));
}

STATIC_INLINE void norm3(double vec[3], double res[3])
{
  double mag = mag3(vec);
  res[0] = vec[0] / mag;
  res[1] = vec[1] / mag;
  res[2] = vec[2] / mag;
}

STATIC_INLINE void add3(double vec1[3], double vec2[3], double res[3])
{
  res[0] = vec1[0] + vec2[0];
  res[1] = vec1[1] + vec2[1];
  res[2] = vec1[2] + vec2[2];
}

STATIC_INLINE void sub3(double vec1[3], double vec2[3], double res[3])
{
  res[0] = vec1[0] - vec2[0];
  res[1] = vec1[1] - vec2[1];
  res[2] = vec1[2] - vec2[2];
}

// Note res cannot be vec1 or vec2
STATIC_INLINE void cross(double vec1[3], double vec2[3], double res[3])
{
  res[0] = vec1[1] * vec2[2] - vec1[2] * vec2[1];
  res[1] = vec1[2] * vec2[0] - vec1[0] * vec2[2];
  res[2] = vec1[0] * vec2[1] - vec1[1] * vec2[0];
}

STATIC_INLINE void scale_vec3(double scl, double vec[3], double res[3])
{
  res[0] = scl * vec[0];
  res[1] = scl * vec[1];
  res[2] = scl * vec[2];
}

STATIC_INLINE bool inside(REAL pos[3])
{
  // XXX - Hack specific to ../../data/vfrhand-nohip.nhdr"
  // but for some reason still not matching what Diderot says
  if(pos[0] < 2*0.9375 || pos[0] > (176-3)*0.9375) return false;
  if(pos[1] < 2*0.9375 || pos[1] > (187-3)*0.9375) return false;
  if(pos[2] < 2.0      || pos[2] > (190-3)) return false;
  return true;
}

STATIC_INLINE double lerp(double out_min, double out_max, double in_min, double in, double in_max)
{
  return (out_min * (in_max - in) + out_max * (in - in_min)) / (in_max - in_min);
}

STATIC_INLINE unsigned char ulerp(double in_min, double in, double in_max)
{
  return (unsigned char)((0.0f * (in_max - in) + 255.0f * (in - in_min)) / (in_max - in_min));
}

STATIC_INLINE REAL
alpha(float *data, REAL val, REAL gmag) {
  unsigned int vi, gi;
  REAL vv, gg, vf, gf;

  /* XXX HEY HACK */
  vv = val/7.8431372549019605 + 3;
  gg = gmag/3.5294117647058822 + 3;
  vi = vv; vf = vv - vi;
  gi = gg; gf = gg - gi;
  return (data[vi+0 + 262*(gi+0)]*(1-vf)*(1-gf) +
          data[vi+1 + 262*(gi+0)]*vf*(1-gf) +
          data[vi+0 + 262*(gi+1)]*(1-vf)*gf +
          data[vi+1 + 262*(gi+1)]*vf*gf);
}

STATIC_INLINE REAL
shade(float *data, REAL val) {
  unsigned int vi;
  REAL vv, vf;

  /* XXX HEY HACK */
  vv = val/0.010050251256281407 + 103;
  vi = vv; vf = vv - vi;
  return (data[vi]*(1-vf) + data[vi+1]*vf);
}

STATIC_INLINE REAL
silho(float *data, REAL val) {
  unsigned int vi;
  REAL vv, vf;

  /* XXX HEY HACK */
  vv = val/0.010050251256281407 + 3;
  vi = vv; vf = vv - vi;
  return (data[vi]*(1-vf) + data[vi+1]*vf);
}

STATIC_INLINE REAL
rival(float *data, REAL rad, REAL the) {
  unsigned int ri, ti;
  REAL rr, tt, rf, tf;

  /* XXX HEY HACK */
  rr = rad/0.0050251256281407036 + 3;
  tt = (the + 3.0*AIR_PI/4.0)/0.015786894472361809 + 3;
  ri = rr; rf = rr - ri;
  ti = tt; tf = tt - ti;
  return (data[ri+0 + 206*(ti+0)]*(1-rf)*(1-tf) +
          data[ri+1 + 206*(ti+0)]*rf*(1-tf) +
          data[ri+0 + 206*(ti+1)]*(1-rf)*tf +
          data[ri+1 + 206*(ti+1)]*rf*tf);
}

STATIC_INLINE void
depth(float *data, REAL rgb[3], REAL dep) {
  unsigned int di;
  REAL dd, df;

  /* XXX HEY HACK */
  dd = dep/0.0050251256281407036 + 3;
  di = dd; df = dd - di;
  ELL_3V_SCALE_ADD2(rgb,
                    1-df, data + 3*di,
                    df, data + 3*(di+1));
}

int main (int argc, const char *argv[])
{
  const char *me = argv[0];
  char *inFile = "../../data/vfrhand-nohip-smooth.nrrd",
    *alphaFile = "../../data/txf/alpha-bone.nrrd",
    *shadeFile = "../../data/txf/shade.nrrd",
    *silhoFile = "../../data/txf/silho.nrrd",
    *rivalFile = "../../data/txf/ridgvall.nrrd",
    *depthFile = "../../data/txf/depth.nrrd";
  float *alphaData = NULL, *shadeData = NULL, *silhoData = NULL,
    *rivalData = NULL, *depthData = NULL;
  char *outfile = "bmark-teem.nrrd";
  double camEye[3] = {127.331, -1322.05, 272.53};
  double camAt[3] = {63.0, 82.6536, 98.0};
  double camUp[3] = {0.9987, 0.0459166, -0.0221267};
  double camNear = -78.0;
  double camFar = 78.0;
  double camFOV = 5.0;

  int imgResU = 640;
  int imgResV = 480;
  double rayStep = 0.15;
  /*
  int imgResU = 480;
  int imgResV = 360;
  double rayStep = 1;
  */

  double temp[3];
  double temp2[3];
  double temp3[3];
  sub3(camAt, camEye, temp);
  double camDist = mag3( temp );
  double camVspNear = camDist + camNear;
  double camVspFar = camDist + camFar;
  double camN[3];
  norm3(temp, camN);
  double camU[3];
  cross(camN, camUp, temp);
  norm3(temp, camU);
  double camV[3];
  cross(camN, camU, camV);
  double camVmax = tan(camFOV*AIR_PI/360.0)*camDist;
  double camUmax = camVmax*(double)(imgResU)/(double)(imgResV);
  double lightVspDir[3] = {-2.0, -3.0, -2.0};
  scale_vec3(lightVspDir[0], camU, temp);
  scale_vec3(lightVspDir[1], camV, temp2);
  add3(temp, temp2, temp);
  scale_vec3(lightVspDir[2], camN, temp2);
  add3(temp, temp2, temp);
  double lightDir[3];
  norm3(temp, lightDir);
  double phongKa = 0.05;
  double phongKd = 0.80;
  double phongKs = 0.20;
  double phongSp = 50.0;
#define sthick 0.4
  double refStep = 1.5;

  // Read in the data
  char* err;
  Nrrd *nin, *nalpha, *nshade, *nsilho, *nrival, *ndepth;
  int status;
  airArray *mop;
  mop = airMopNew();

  nin = nrrdNew();
  airMopAdd(mop, nin, (airMopper)nrrdNuke, airMopAlways);
  nalpha = nrrdNew();
  airMopAdd(mop, nalpha, (airMopper)nrrdNuke, airMopAlways);
  nshade = nrrdNew();
  airMopAdd(mop, nshade, (airMopper)nrrdNuke, airMopAlways);
  nsilho = nrrdNew();
  airMopAdd(mop, nsilho, (airMopper)nrrdNuke, airMopAlways);
  nrival = nrrdNew();
  airMopAdd(mop, nrival, (airMopper)nrrdNuke, airMopAlways);
  ndepth = nrrdNew();
  airMopAdd(mop, ndepth, (airMopper)nrrdNuke, airMopAlways);

  if (nrrdLoad(nin, inFile, NULL)
      || nrrdLoad(nalpha, alphaFile, NULL)
      || nrrdLoad(nshade, shadeFile, NULL)
      || nrrdLoad(nsilho, silhoFile, NULL)
      || nrrdLoad(nrival, rivalFile, NULL)
      || nrrdLoad(ndepth, depthFile, NULL)) {
    airMopAdd(mop, err = biffGetDone(NRRD), airFree, airMopAlways);
    fprintf(stderr, "%s: Trouble reading:\n%s", me, err);
    airMopError(mop); 
    return -1;
  }
  alphaData = AIR_CAST(float *, nalpha->data); // XXX HEY HACK
  shadeData = AIR_CAST(float *, nshade->data); // XXX HEY HACK
  silhoData = AIR_CAST(float *, nsilho->data); // XXX HEY HACK
  rivalData = AIR_CAST(float *, nrival->data); // XXX HEY HACK
  depthData = AIR_CAST(float *, ndepth->data); // XXX HEY HACK

  gageContext *ctx;
  gagePerVolume *pvl;
  const double *valP, *gradP, *hessP;
  double kparm[NRRD_KERNEL_PARMS_NUM] = {0.0};

  ctx = gageContextNew();
  airMopAdd(mop, ctx, (airMopper)gageContextNix, airMopAlways);
  if (!( pvl = gagePerVolumeNew(ctx, nin, gageKindScl) )) {
    airMopAdd(mop, err = biffGetDone(GAGE), airFree, airMopAlways);
    fprintf(stderr, "%s: couldn't create gage context:\n%s", me, err);
    airMopError(mop); 
    return -1;
  }
  if (gagePerVolumeAttach(ctx, pvl)
      || gageKernelSet(ctx, gageKernel00, nrrdKernelBSpline3, kparm)
      || gageKernelSet(ctx, gageKernel11, nrrdKernelBSpline3D, kparm)
      || gageKernelSet(ctx, gageKernel22, nrrdKernelBSpline3DD, kparm)
      || gageQueryItemOn(ctx, pvl, gageSclValue)
      || gageQueryItemOn(ctx, pvl, gageSclGradVec)
      || gageQueryItemOn(ctx, pvl, gageSclHessian)
      || gageUpdate(ctx)) {
    airMopAdd(mop, err = biffGetDone(GAGE), airFree, airMopAlways);
    fprintf(stderr, "%s: Trouble setting up gage context:\n%s", me, err);
    airMopError(mop); 
    return -1;
  }

  if (!( (valP = gageAnswerPointer(ctx, pvl, gageSclValue))
         && (gradP = gageAnswerPointer(ctx, pvl, gageSclGradVec))
         && (hessP = gageAnswerPointer(ctx, pvl, gageSclHessian)) )) {
    fprintf(stderr, "%s: Trouble getting answer pointer\n", me);
    airMopError(mop); 
    return -1;
  }

  int ui, vi;
  REAL rayU, rayV, rayN;
  REAL rayVec[3];
  REAL vv[3];
  REAL rayPos[3];
  REAL rayTransp, rayRGB[3];
  REAL norm[3];
  REAL ld;
  REAL hd;
  REAL mat;

  // allocate output image
  Nrrd *nimg = nrrdNew();
  airMopAdd(mop, nimg, (airMopper)nrrdNuke, airMopAlways);
  if (nrrdAlloc_va(nimg, nrrdTypeDouble, 3,
                   AIR_CAST(size_t, 4),
                   AIR_CAST(size_t, imgResU),
                   AIR_CAST(size_t, imgResV))) {
    airMopAdd(mop, err = biffGetDone(NRRD), airFree, airMopAlways);
    fprintf(stderr, "%s: Trouble allocating image output:\n%s", me, err);
    airMopError(mop); 
    return -1;
  }
  double *out_data = AIR_CAST(double *, nimg->data);
  /*
  fprintf(stderr, "%s: hit return\n", me);
  fgetc(stdin);
  fprintf(stderr, "%s: here we go\n", me);
  */
  double t0 = airTime(); // start timing
  for (vi = 0; vi < imgResV; vi++) {
    rayV = lerp(-camVmax, camVmax, -0.5, (double)(vi), (double)(imgResV) - 0.5);
    for (ui = 0; ui < imgResU; ui++) {
      REAL len;
      rayU = lerp(-camUmax, camUmax, -0.5, (double)(ui), (double)(imgResU) - 0.5);
      ELL_3V_SCALE_ADD3(rayVec,   1.0, camN,   rayU/camDist, camU,   rayV/camDist, camV);
      ELL_3V_SCALE(vv, -1, rayVec);
      ELL_3V_NORM(vv, vv, len);
      rayTransp = 1;
      ELL_3V_SET(rayRGB, 0, 0, 0);

      for (rayN = camVspNear; rayN <= camVspFar; rayN += rayStep) {
        ELL_3V_SCALE_ADD2(rayPos,   1.0, camEye,  rayN, rayVec);

        if (inside(rayPos)) {
          if (gageProbeSpace(ctx, rayPos[0], rayPos[1], rayPos[2],
                             AIR_FALSE /* index space */, AIR_TRUE /* clamp */)) {
            fprintf(stderr, "%s: Trouble with gageProbe:\n%s", me, ctx->errStr);
            airMopError(mop); return -1;
          }
          REAL gmag = ELL_3V_LEN(gradP);
          REAL aa = alpha(alphaData, AIR_CLAMP(0, *valP, 2000), AIR_CLAMP(0, gmag, 900));
          REAL matRGB[3];
          if (aa > 0.0) {
            REAL norm[3], gray, gten[9], proj[9], matA[9], matB[9];
            ELL_3V_SCALE(norm, -1.0/gmag, gradP); 
            aa = 1.0 - pow(1.0-aa, rayStep/refStep);
            gray = shade(shadeData, ELL_3V_DOT(norm, lightDir));
            ELL_3MV_OUTER(matA, norm, norm);
            ELL_3M_SCALE(proj, -1, matA);
            proj[0] += 1; proj[4] += 1; proj[8] += 1;

            ELL_3M_MUL(matB, hessP, proj);
            ELL_3M_MUL(gten, proj, matB);
            ELL_3M_SCALE(gten, -1.0/gmag, gten);
            REAL kv = ELL_3MV_CONTR(gten, vv)/ELL_3MV_CONTR(proj, vv);
            kv = AIR_CLAMP(0.0, kv, 1.0/sthick);
            REAL vdn = ELL_3V_DOT(vv, norm);
            vdn = vdn*(1.0 + 2.0*vdn*vdn);
            REAL tmp = sthick*(kv - 1.0/sthick);
            REAL sarg = vdn*vdn + tmp*tmp;
            gray *= silho(silhoData, AIR_CLAMP(0, sarg, 2));

            REAL gtns = ELL_3M_FROBSQRD(gten);
            REAL gttr = ELL_3M_TRACE(gten);
            REAL disc = SQRT(2.0*gtns - gttr*gttr);
            REAL k1 = (gttr + disc)/2.0;
            REAL k2 = (gttr - disc)/2.0;
            tmp = SQRT(gtns);
            gray *= rival(rivalData, AIR_MIN(0.55, tmp)/0.55, ATAN2(k2,k1));

            depth(depthData, matRGB, lerp(0.0, 1.0, camVspNear, rayN, camVspFar));

            ELL_3V_SCALE_INCR(rayRGB, rayTransp*aa*gray, matRGB);
            rayTransp *= 1 - aa;
          }
        }
        if (rayTransp < 0.01) {
          rayTransp = 0;
          break;
        }
      }

      ELL_4V_SET(out_data + 4*(ui + imgResU*vi), 
                 0.9*rayRGB[0], 0.9*rayRGB[1], 0.9*rayRGB[2], 1.0-rayTransp);

    }
  }

  double totalTime = airTime() - t0;
  printf("usr=%f\n", totalTime);

  if (nrrdSave(outfile, nimg, NULL)) {
    airMopAdd(mop, err = biffGetDone(NRRD), airFree, airMopAlways);
    fprintf(stderr, "%s: Trouble writing output:\n%s", me, err);
    airMopError(mop); 
    return -1;
  }
  
  airMopOkay(mop); 
  return 0;

}