// Geometry gradients for pathtracer-diff, direct-lighting scope.
//
// Interior term (k_geo_interior): the derivative of the direct-lighting
// image with respect to vertex positions under DETACHED sampling: camera
// ray directions and light-point barycentrics are held fixed while the
// receiver hit point, normals, distances, cosines, and the light-area
// measure move with the vertices. Derivatives are propagated with
// forward-mode dual numbers, one pass per perturbed coordinate (18 per
// sample: 3 receiver verts + 3 light verts, xyz), and scattered into
// grad_verts[V, 3] with atomics. Diffuse receivers and emitters only;
// uv motion across textures is not differentiated.
//
// Boundary terms (edge sampling, Li et al. 2018 "redner"):
// - k_geo_boundary: shadow silhouettes. Sample a scene edge point p and a
//   light point y, cast y->p onward to the receiver x_b, test that the
//   edge is a silhouette from y, and accumulate (radiance jump) x
//   (projective sweep velocity of the shadow boundary through x_b) x
//   (receiver-area -> pixel measure) x dLoss/dpixel into the edge's two
//   vertices.
// - k_geo_primary: camera silhouettes. Sample an edge point, require it to
//   be a silhouette from the camera and directly visible, and accumulate
//   (front-minus-background direct radiance) x (image-space sweep velocity
//   of the projected edge) x dLoss/dpixel.
//
// Together: grad_verts = interior (smooth) + shadow boundary + primary
// boundary for the direct-lighting transport term. Indirect bounces are
// not differentiated with respect to geometry.

#include <cuda_runtime.h>
#include <curand_kernel.h>

#include "geometry_grad.h"

namespace {

constexpr int kThreads = 128;
constexpr float kPi = 3.14159265358979323846f;
constexpr float kInvPi = 0.31830988618379067154f;
constexpr float kRayEps = 1e-4f;
constexpr float kShadowEps = 1e-3f;

// ------------------------------------------------------------ small vec ops
struct V3 {
  float x, y, z;
};
__device__ __forceinline__ V3 v3(float x, float y, float z) {
  V3 r{x, y, z};
  return r;
}
__device__ __forceinline__ V3 operator+(V3 a, V3 b) {
  return v3(a.x + b.x, a.y + b.y, a.z + b.z);
}
__device__ __forceinline__ V3 operator-(V3 a, V3 b) {
  return v3(a.x - b.x, a.y - b.y, a.z - b.z);
}
__device__ __forceinline__ V3 operator*(V3 a, float s) {
  return v3(a.x * s, a.y * s, a.z * s);
}
__device__ __forceinline__ float vdot(V3 a, V3 b) {
  return a.x * b.x + a.y * b.y + a.z * b.z;
}
__device__ __forceinline__ V3 vcross(V3 a, V3 b) {
  return v3(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z,
            a.x * b.y - a.y * b.x);
}
__device__ __forceinline__ V3 vnorm(V3 a) {
  float s = rsqrtf(fmaxf(vdot(a, a), 1e-30f));
  return a * s;
}

// ----------------------------------------------------------- dual numbers
struct D {
  float v, d;
};
__device__ __forceinline__ D dc(float v) { D r{v, 0.0f}; return r; }
__device__ __forceinline__ D dv(float v, float d) { D r{v, d}; return r; }
__device__ __forceinline__ D operator+(D a, D b) { return dv(a.v + b.v, a.d + b.d); }
__device__ __forceinline__ D operator-(D a, D b) { return dv(a.v - b.v, a.d - b.d); }
__device__ __forceinline__ D operator*(D a, D b) {
  return dv(a.v * b.v, a.v * b.d + a.d * b.v);
}
__device__ __forceinline__ D operator/(D a, D b) {
  float inv = 1.0f / b.v;
  return dv(a.v * inv, (a.d - a.v * b.d * inv) * inv);
}
__device__ __forceinline__ D dsqrt(D a) {
  float s = sqrtf(fmaxf(a.v, 1e-20f));
  return dv(s, 0.5f * a.d / s);
}
__device__ __forceinline__ D dmax0(D a) {
  return a.v > 0.0f ? a : dc(0.0f);
}

struct DV3 {
  D x, y, z;
};
__device__ __forceinline__ DV3 dv3(D x, D y, D z) { DV3 r{x, y, z}; return r; }
__device__ __forceinline__ DV3 dvc(V3 a) {
  return dv3(dc(a.x), dc(a.y), dc(a.z));
}
__device__ __forceinline__ DV3 operator+(DV3 a, DV3 b) {
  return dv3(a.x + b.x, a.y + b.y, a.z + b.z);
}
__device__ __forceinline__ DV3 operator-(DV3 a, DV3 b) {
  return dv3(a.x - b.x, a.y - b.y, a.z - b.z);
}
__device__ __forceinline__ DV3 operator*(DV3 a, D s) {
  return dv3(a.x * s, a.y * s, a.z * s);
}
__device__ __forceinline__ D ddot(DV3 a, DV3 b) {
  return a.x * b.x + a.y * b.y + a.z * b.z;
}
__device__ __forceinline__ DV3 dcross(DV3 a, DV3 b) {
  return dv3(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z,
             a.x * b.y - a.y * b.x);
}

// ---------------------------------------------------------- scene helpers
__device__ __forceinline__ V3 vert(const PtdSceneArgs& a, int face, int corner) {
  const float* t = &a.tris[face * 9 + corner * 3];
  return v3(t[0], t[1], t[2]);
}

// closest-hit over the BVH (values only; reuses main-kernel node layout)
__device__ int geo_closest(const PtdSceneArgs& a, V3 ro, V3 rd, float tmin,
                           float& tbest, float& bu, float& bv) {
  float roa[3] = {ro.x, ro.y, ro.z};
  float dira[3] = {rd.x, rd.y, rd.z};
  float inv[3];
  {
    float d;
    d = rd.x; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[0] = 1.0f / d;
    d = rd.y; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[1] = 1.0f / d;
    d = rd.z; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[2] = 1.0f / d;
  }
  int stack[64];
  int sp = 0;
  stack[sp++] = 0;
  int best = -1;
  while (sp > 0) {
    int nid = stack[--sp];
    const float* b = &a.nodes_f[nid * 6];
    float t0 = tmin, t1 = tbest;
    bool hit = true;
#pragma unroll
    for (int ax = 0; ax < 3; ++ax) {
      float lo = (b[ax] - roa[ax]) * inv[ax];
      float hi = (b[3 + ax] - roa[ax]) * inv[ax];
      if (lo > hi) { float tmp = lo; lo = hi; hi = tmp; }
      t0 = fmaxf(t0, lo);
      t1 = fminf(t1, hi);
    }
    hit = t0 <= t1;
    if (!hit) continue;
    const int* n = &a.nodes_i[nid * 3];
    if (n[2] & 1) {
      for (int f = n[0]; f < n[0] + n[1]; ++f) {
        const float* tv = &a.tris[f * 9];
        V3 v0 = v3(tv[0], tv[1], tv[2]);
        V3 e1 = v3(tv[3], tv[4], tv[5]) - v0;
        V3 e2 = v3(tv[6], tv[7], tv[8]) - v0;
        V3 p = vcross(rd, e2);
        float det = vdot(e1, p);
        if (fabsf(det) < 1e-12f) continue;
        float invd = 1.0f / det;
        V3 s = ro - v0;
        float u = vdot(s, p) * invd;
        if (u < -1e-6f || u > 1.0f + 1e-6f) continue;
        V3 q = vcross(s, e1);
        float w = vdot(rd, q) * invd;
        if (w < -1e-6f || u + w > 1.0f + 1e-6f) continue;
        float tt = vdot(e2, q) * invd;
        if (tt < tmin || tt > tbest) continue;
        tbest = tt;
        bu = u;
        bv = w;
        best = f;
      }
    } else if (sp + 2 <= 64) {
      int axis = n[2] >> 1;
      int nearc = (dira[axis] >= 0.0f) ? n[0] : n[1];
      int farc = (dira[axis] >= 0.0f) ? n[1] : n[0];
      stack[sp++] = farc;
      stack[sp++] = nearc;
    }
  }
  return best;
}

__device__ bool geo_occluded(const PtdSceneArgs& a, V3 ro, V3 rd, float tmax) {
  float tb = tmax;
  float bu, bv;
  int f = geo_closest(a, ro, rd, kRayEps, tb, bu, bv);
  return f >= 0;
}

__device__ __forceinline__ int cdf_pick_geo(const float* cdf, int n, float r) {
  int lo = 0, hi = n - 1;
  while (lo < hi) {
    int mid = (lo + hi) >> 1;
    if (cdf[mid] < r) lo = mid + 1; else hi = mid;
  }
  return lo;
}

// constant-per-material fetches (geometry pass ignores texture footprints;
// values are taken at the frozen uv, texel 0 approximation for 1x1 and the
// bilinear value otherwise is close enough for the smooth term's weights)
__device__ __forceinline__ void mat_albedo0(const PtdSceneArgs& a, int m,
                                            float rgb[3]) {
  const int* h = &a.tex_hdr[m * 3];
  const float* t = &a.tex[h[0] * 3];
  rgb[0] = t[0];
  rgb[1] = t[1];
  rgb[2] = t[2];
}
__device__ __forceinline__ void mat_emission0(const PtdSceneArgs& a, int m,
                                              float rgb[3]) {
  const int* h = &a.emi_hdr[m * 3];
  const float* t = &a.emi_tex[h[0] * 3];
  rgb[0] = t[0];
  rgb[1] = t[1];
  rgb[2] = t[2];
}

// The direct-lighting NEE term as a dual-number function of one perturbed
// vertex coordinate. Inputs: receiver verts r0..r2 (duals), light verts
// l0..l2 (duals), camera origin o and FIXED direction w (detached), FIXED
// light barycentrics (b0, b1). Returns f_geo = cos_x * cos_y / d2 *
// |cross(le1, le2)| (the area-measure-weighted geometry term including the
// light-area jacobian; the caller multiplies by albedo/pi * E * 2 /
// (2 * pdf-const)).
__device__ D direct_geo_dual(DV3 r0, DV3 r1, DV3 r2, DV3 l0, DV3 l1, DV3 l2,
                             V3 o, V3 w, float b0, float b1) {
  // receiver hit: t = dot(r0 - o, ng) / dot(w, ng)
  DV3 e1 = r1 - r0;
  DV3 e2 = r2 - r0;
  DV3 ng = dcross(e1, e2);
  DV3 ov = dvc(o);
  DV3 wv = dvc(w);
  D num = ddot(r0 - ov, ng);
  D den = ddot(wv, ng);
  D t = num / den;
  DV3 x = ov + wv * t;
  D nglen = dsqrt(ddot(ng, ng));
  // light point and its (unnormalized) normal
  DV3 le1 = l1 - l0;
  DV3 le2 = l2 - l0;
  DV3 lng = dcross(le1, le2);
  D lnglen = dsqrt(ddot(lng, lng));
  DV3 y = l0 + le1 * dc(b0) + le2 * dc(b1);
  DV3 dvec = y - x;
  D d2 = ddot(dvec, dvec);
  D d = dsqrt(d2);
  // cos_x = |dot(ng, dvec)| / (|ng| d); cos_y = |dot(lng, dvec)| / (|lng| d)
  D cx = ddot(ng, dvec) / (nglen * d);
  cx = dv(fabsf(cx.v), cx.v >= 0.0f ? cx.d : -cx.d);
  D cy = ddot(lng, dvec) / (lnglen * d);
  cy = dv(fabsf(cy.v), cy.v >= 0.0f ? cy.d : -cy.d);
  // per-face area measure: dA = |lng| / 2; the light list samples the face
  // uniformly by area, so the estimator carries |lng|/2 / (frozen pdf-area)
  D geo = cx * cy / d2 * (lnglen * dc(0.5f));
  return dmax0(geo);
}

struct GeoCam {
  float c[12];
};

}  // namespace

// ------------------------------------------------------- interior kernel
namespace {
__global__ void k_geo_interior(PtdSceneArgs a, const int* face_verts,
                               GeoCam gc, int H, int W, int spp,
                               unsigned long long seed,
                               const float* grad_image, float* grad_verts) {
  const float* cam = gc.c;
  int pid = blockIdx.x * blockDim.x + threadIdx.x;
  if (pid >= H * W) return;
  int px = pid % W, py = pid / W;
  V3 o = v3(cam[0], cam[1], cam[2]);
  float gpix[3] = {grad_image[pid * 3], grad_image[pid * 3 + 1],
                   grad_image[pid * 3 + 2]};
  if (gpix[0] == 0.0f && gpix[1] == 0.0f && gpix[2] == 0.0f) return;
  float inv_spp = 1.0f / (float)spp;

  for (int s = 0; s < spp; ++s) {
    curandStatePhilox4_32_10_t st;
    curand_init(seed, (unsigned long long)pid * spp + s, 0, &st);
    float jx = curand_uniform(&st);
    float jy = curand_uniform(&st);
    float nx = 2.0f * ((px + jx) / (float)W) - 1.0f;
    float ny = 1.0f - 2.0f * ((py + jy) / (float)H);
    V3 w = vnorm(v3(cam[3] + nx * cam[6] + ny * cam[9],
                    cam[4] + nx * cam[7] + ny * cam[10],
                    cam[5] + nx * cam[8] + ny * cam[11]));
    float tb = 1e30f, bu, bv;
    int face = geo_closest(a, o, w, kRayEps, tb, bu, bv);
    if (face < 0 || a.n_lights <= 0) continue;
    int m = a.mat_ids[face];
    if (a.mat_type[m] != 0) continue;  // diffuse receivers only
    // light sample (detached picks)
    float r1 = curand_uniform(&st);
    float r2 = curand_uniform(&st);
    float r3 = curand_uniform(&st);
    int li = 0;
    {
      int lo = 0, hi = a.n_lights - 1;
      while (lo < hi) {
        int mid = (lo + hi) >> 1;
        if (a.light_cdf[mid] < r1) lo = mid + 1; else hi = mid;
      }
      li = lo;
    }
    int lface = a.light_faces[li];
    float su = sqrtf(r2);
    float b0 = 1.0f - su, b1 = r3 * su;
    // occlusion at the primal configuration (frozen for the interior term)
    V3 rv[3] = {vert(a, face, 0), vert(a, face, 1), vert(a, face, 2)};
    V3 lv[3] = {vert(a, lface, 0), vert(a, lface, 1), vert(a, lface, 2)};
    V3 x = o + w * tb;
    V3 y = lv[0] + (lv[1] - lv[0]) * b0 + (lv[2] - lv[0]) * b1;
    V3 dl = y - x;
    float dist = sqrtf(fmaxf(vdot(dl, dl), 1e-12f));
    V3 wi = dl * (1.0f / dist);
    V3 ngp = vcross(rv[1] - rv[0], rv[2] - rv[0]);
    V3 np = vnorm(ngp);
    if (vdot(np, w) > 0.0f) np = np * -1.0f;
    if (vdot(np, wi) <= 1e-6f) continue;
    if (geo_occluded(a, x + np * kRayEps, wi, dist - kShadowEps)) continue;

    float alb[3], em[3];
    mat_albedo0(a, m, alb);
    mat_emission0(a, a.mat_ids[lface], em);
    // frozen-pdf estimator constant: (1/pi) * E * (2 / (area_frac)) where
    // the face was picked with prob area_face/area_total and sampled at
    // density 2/|lng|; combined constant = larea_total/pi times the dual
    // geo term normalized by the face area fraction. The dual geo term
    // already carries |lng|/2, so multiply by larea / area_face:
    V3 lngp = vcross(lv[1] - lv[0], lv[2] - lv[0]);
    float area_face = 0.5f * sqrtf(fmaxf(vdot(lngp, lngp), 1e-20f));
    float scale_c = kInvPi * a.total_light_area / fmaxf(area_face, 1e-12f) *
                    inv_spp;

    int rvi[3] = {face_verts[face * 3], face_verts[face * 3 + 1],
                  face_verts[face * 3 + 2]};
    int lvi[3] = {face_verts[lface * 3], face_verts[lface * 3 + 1],
                  face_verts[lface * 3 + 2]};

    // 18 dual passes: receiver verts then light verts, xyz each
    for (int vi = 0; vi < 6; ++vi) {
      for (int cc = 0; cc < 3; ++cc) {
        DV3 R[3] = {dvc(rv[0]), dvc(rv[1]), dvc(rv[2])};
        DV3 L[3] = {dvc(lv[0]), dvc(lv[1]), dvc(lv[2])};
        D* slot;
        if (vi < 3)
          slot = (cc == 0 ? &R[vi].x : cc == 1 ? &R[vi].y : &R[vi].z);
        else
          slot = (cc == 0 ? &L[vi - 3].x : cc == 1 ? &L[vi - 3].y
                                                   : &L[vi - 3].z);
        slot->d = 1.0f;
        D geo = direct_geo_dual(R[0], R[1], R[2], L[0], L[1], L[2], o, w,
                                b0, b1);
        if (geo.d == 0.0f) continue;
        float gsum = 0.0f;
        for (int c = 0; c < 3; ++c)
          gsum += gpix[c] * alb[c] * em[c];
        float contrib = gsum * scale_c * geo.d;
        int target = (vi < 3 ? rvi[vi] : lvi[vi - 3]);
        atomicAdd(&grad_verts[target * 3 + cc], contrib);
      }
    }
  }
}

// ------------------------------------------------------- boundary kernel
__global__ void k_geo_boundary(PtdSceneArgs a, const int* face_verts,
                               const int* edges, int n_edges,
                               const float* edge_cdf, GeoCam gc, int H,
                               int W, int n_samples, unsigned long long seed,
                               const float* grad_image, float* grad_verts,
                               float total_edge_len) {
  const float* cam = gc.c;
  int sid = blockIdx.x * blockDim.x + threadIdx.x;
  if (sid >= n_samples || a.n_lights <= 0) return;
  curandStatePhilox4_32_10_t st;
  curand_init(seed ^ 0x9e3779b97f4a7c15ull, (unsigned long long)sid, 0, &st);
  float r0 = curand_uniform(&st);
  float rs = curand_uniform(&st);
  float r1 = curand_uniform(&st);
  float r2 = curand_uniform(&st);
  float r3 = curand_uniform(&st);
  // pick an edge by length, a point on it, a light point
  int ei;
  {
    int lo = 0, hi = n_edges - 1;
    while (lo < hi) {
      int mid = (lo + hi) >> 1;
      if (edge_cdf[mid] < r0) lo = mid + 1; else hi = mid;
    }
    ei = lo;
  }
  // edge row: (corner_a, corner_b, face_a, face_b); corners index into
  // face_a, face_b = -1 for border edges
  const int* E = &edges[ei * 4];
  int fa = E[2];
  int fb = E[3];
  int ca = E[0] & 3;
  int cb = E[1] & 3;
  int va = face_verts[fa * 3 + ca];
  int vb = face_verts[fa * 3 + cb];
  V3 p0 = vert(a, fa, ca);
  V3 p1 = vert(a, fa, cb);
  V3 pe = p0 + (p1 - p0) * rs;

  // light point
  int li;
  {
    int lo = 0, hi = a.n_lights - 1;
    while (lo < hi) {
      int mid = (lo + hi) >> 1;
      if (a.light_cdf[mid] < r1) lo = mid + 1; else hi = mid;
    }
    li = lo;
  }
  int lface = a.light_faces[li];
  V3 l0 = vert(a, lface, 0), l1 = vert(a, lface, 1), l2 = vert(a, lface, 2);
  float su = sqrtf(r2);
  V3 y = l0 + (l1 - l0) * (1.0f - su) + (l2 - l0) * (r3 * su);

  // silhouette test from y: adjacent faces straddle the plane through y
  V3 nga = vcross(vert(a, fa, 1) - vert(a, fa, 0),
                  vert(a, fa, 2) - vert(a, fa, 0));
  float sa_side = vdot(nga, y - vert(a, fa, 0));
  bool sil;
  if (fb < 0) {
    sil = true;  // border edge: always a silhouette
  } else {
    V3 ngb = vcross(vert(a, fb, 1) - vert(a, fb, 0),
                    vert(a, fb, 2) - vert(a, fb, 0));
    float sb_side = vdot(ngb, y - vert(a, fb, 0));
    sil = (sa_side > 0.0f) != (sb_side > 0.0f);
  }
  if (!sil) return;

  // cast y -> pe onward to the receiver
  V3 wdir = pe - y;
  float dpe = sqrtf(fmaxf(vdot(wdir, wdir), 1e-12f));
  wdir = wdir * (1.0f / dpe);
  float tb = 1e30f, bu, bv;
  int rface = geo_closest(a, y, wdir, dpe + kShadowEps, tb, bu, bv);
  if (rface < 0) return;
  int rm = a.mat_ids[rface];
  if (a.mat_type[rm] != 0) return;  // diffuse receivers only
  V3 xb = y + wdir * tb;

  // receiver must be visible from the camera; find its pixel
  V3 o = v3(cam[0], cam[1], cam[2]);
  V3 toc = xb - o;
  float dcam = sqrtf(fmaxf(vdot(toc, toc), 1e-12f));
  V3 wc = toc * (1.0f / dcam);
  {
    float tb2 = dcam - kShadowEps;
    float b2u, b2v;
    int f2 = geo_closest(a, o, wc, kRayEps, tb2, b2u, b2v);
    if (f2 >= 0) return;  // camera-occluded
  }
  // pixel coordinates: invert the camera basis (fwd, rs, us are rows 3..11)
  V3 fwd = v3(cam[3], cam[4], cam[5]);
  V3 rsv = v3(cam[6], cam[7], cam[8]);
  V3 usv = v3(cam[9], cam[10], cam[11]);
  float rs2 = vdot(rsv, rsv), us2 = vdot(usv, usv);
  float wf = vdot(wc, fwd) / vdot(fwd, fwd);
  if (wf <= 1e-6f) return;
  V3 wplane = wc * (1.0f / (wf * vdot(fwd, fwd)));
  // wc/(w.f/|f|^2) = fwd + nx*rs + ny*us  -> project
  float nx = vdot(wplane - fwd, rsv) / rs2;
  float ny = vdot(wplane - fwd, usv) / us2;
  int pxi = (int)((nx + 1.0f) * 0.5f * W);
  int pyi = (int)((1.0f - ny) * 0.5f * H);
  if (pxi < 0 || pxi >= W || pyi < 0 || pyi >= H) return;
  const float* gpix = &grad_image[(pyi * W + pxi) * 3];

  // radiance jump: the direct term from y at xb (present on the unblocked
  // side of the sweeping shadow boundary)
  V3 ngr = vcross(vert(a, rface, 1) - vert(a, rface, 0),
                  vert(a, rface, 2) - vert(a, rface, 0));
  V3 nr = vnorm(ngr);
  if (vdot(nr, wdir) > 0.0f) nr = nr * -1.0f;
  float cxr = fmaxf(vdot(nr, wdir * -1.0f), 1e-6f);
  V3 lng = vcross(l1 - l0, l2 - l0);
  V3 ln = vnorm(lng);
  if (vdot(ln, xb - y) < 0.0f) ln = ln * -1.0f;
  float cyl = fmaxf(vdot(ln, wdir), 1e-6f);
  float d2 = tb * tb;
  float alb[3], em[3];
  mat_albedo0(a, rm, alb);
  mat_emission0(a, a.mat_ids[lface], em);
  // integrand jump per unit light area, times the light MC factor A_total
  float Lterm = cxr * cyl / d2 * a.total_light_area * kInvPi;

  // the shadow boundary on the receiver is the projective image of the
  // edge from y: P(dp) = kappa * (dp - wdir * <nr, dp> / <nr, wdir>) maps
  // edge-point motion dp to boundary-point motion in the receiver plane.
  // tau = P(unit edge tangent) gives the curve tangent (its length is the
  // edge->curve arc stretch); n_c = nr x tau-hat is the in-plane curve
  // normal.
  float kappa = tb / dpe;
  float nw = vdot(nr, wdir);
  if (fabsf(nw) < 1e-4f) return;  // grazing receiver
  V3 edir = vnorm(p1 - p0);
  V3 tau = (edir - wdir * (vdot(nr, edir) / nw)) * kappa;
  float taul = sqrtf(fmaxf(vdot(tau, tau), 1e-12f));
  V3 nc = vcross(nr, tau * (1.0f / taul));
  float ncl = sqrtf(fmaxf(vdot(nc, nc), 1e-12f));
  nc = nc * (1.0f / ncl);

  // sign: the boundary flux is (jump) * <n_c, v> with n_c pointing INTO
  // the blocked region. Probe BOTH sides at a distance-scaled step and
  // require them to disagree; a boundary point probed at the ray-epsilon
  // scale gives a coin-flip verdict and the contributions cancel to zero.
  float sign;
  {
    float step = fmaxf(0.02f, 0.02f * tb);
    bool blk[2];
    for (int sdx = 0; sdx < 2; ++sdx) {
      V3 probe = xb + nc * (sdx == 0 ? step : -step);
      V3 dl2 = y - probe;
      float dist2 = sqrtf(fmaxf(vdot(dl2, dl2), 1e-12f));
      V3 wi2 = dl2 * (1.0f / dist2);
      blk[sdx] = geo_occluded(a, probe + nr * kRayEps, wi2,
                              dist2 - kShadowEps);
    }
    if (blk[0] == blk[1]) return;  // not a clean boundary crossing here
    sign = blk[0] ? 1.0f : -1.0f;
  }

  // receiver-area -> pixel-mean measure at xb: how many pixel means change
  // per unit swept receiver area. dOmega/dndc^2 = |r||s||f|/|u|^3 with
  // u = f + nx r + ny s and |u| = |f|^2 / <wc, f>; dA = dOmega d^2 / cos;
  // one pixel covers (2/W)(2/H) of ndc.
  float ff = vdot(fwd, fwd);
  float wfr = fmaxf(vdot(wc, fwd), 1e-6f);
  float ul = ff / wfr;
  float cos_cam = fmaxf(vdot(nr, wc * -1.0f), 1e-6f);
  float rho = 0.25f * (float)W * (float)H * cos_cam /
              fmaxf(dcam * dcam, 1e-12f) * ul * ul * ul /
              fmaxf(sqrtf(rs2) * sqrtf(us2) * sqrtf(ff), 1e-12f);

  // pdf: edge by length (cdf ~ length, point ~ 1/length, so the estimator
  // carries total_edge_len), light by area (A_total already in Lterm)
  float base = 0.0f;
  for (int c = 0; c < 3; ++c) base += gpix[c] * alb[c] * em[c];
  base *= Lterm * sign * taul * rho * total_edge_len / (float)n_samples;

  // velocity of xb per unit vertex coordinate: dp = wgt * e_cc through the
  // projective map, dotted with the curve normal
  for (int vi = 0; vi < 2; ++vi) {
    float wgt = (vi == 0) ? (1.0f - rs) : rs;
    int target = (vi == 0) ? va : vb;
    for (int cc = 0; cc < 3; ++cc) {
      V3 dp = v3(cc == 0 ? 1.0f : 0.0f, cc == 1 ? 1.0f : 0.0f,
                 cc == 2 ? 1.0f : 0.0f) * wgt;
      V3 dxb = (dp - wdir * (vdot(nr, dp) / nw)) * kappa;
      float vel = vdot(nc, dxb);
      if (vel != 0.0f)
        atomicAdd(&grad_verts[target * 3 + cc], base * vel);
    }
  }
}

// ------------------------------------------------------- primary kernel
// Primary (camera) silhouettes: the image-space discontinuity where an
// object's outline sweeps across pixels. Sample an edge point, require it
// to be a silhouette from the camera and directly visible, and accumulate
// (front radiance - background radiance) x (image-space sweep velocity of
// the projected edge) x dLoss/dpixel. Radiances are the direct-lighting
// values (emission + one-light-sample diffuse reflection); indirect
// bounces are out of scope, as everywhere in the geometry pass.
__device__ bool geo_direct_radiance(const PtdSceneArgs& a, int face, V3 x,
                                    V3 view_from, float r1, float r2,
                                    float r3, float L[3]) {
  L[0] = L[1] = L[2] = 0.0f;
  if (face < 0) return true;  // miss: no environment in the geometry pass
  int m = a.mat_ids[face];
  float em[3];
  mat_emission0(a, m, em);
  L[0] = em[0];
  L[1] = em[1];
  L[2] = em[2];
  if (a.mat_type[m] != 0) {
    // non-diffuse reflection is out of scope; emitters still radiate
    return em[0] > 0.0f || em[1] > 0.0f || em[2] > 0.0f;
  }
  if (a.n_lights <= 0) return true;
  int li = cdf_pick_geo(a.light_cdf, a.n_lights, r1);
  int lface = a.light_faces[li];
  V3 l0 = vert(a, lface, 0), l1 = vert(a, lface, 1), l2 = vert(a, lface, 2);
  float su = sqrtf(r2);
  V3 y = l0 + (l1 - l0) * (1.0f - su) + (l2 - l0) * (r3 * su);
  V3 ng = vcross(vert(a, face, 1) - vert(a, face, 0),
                 vert(a, face, 2) - vert(a, face, 0));
  V3 n = vnorm(ng);
  if (vdot(n, view_from - x) < 0.0f) n = n * -1.0f;  // face the camera
  V3 dl = y - x;
  float d2 = fmaxf(vdot(dl, dl), 1e-12f);
  float d = sqrtf(d2);
  V3 wi = dl * (1.0f / d);
  float cx = vdot(n, wi);
  if (cx <= 1e-6f) return true;  // light behind: reflected term is zero
  V3 lng = vcross(l1 - l0, l2 - l0);
  V3 ln = vnorm(lng);
  if (vdot(ln, dl) > 0.0f) ln = ln * -1.0f;
  float cy = -vdot(ln, wi);
  if (cy <= 1e-6f) return true;
  if (geo_occluded(a, x + n * kRayEps, wi, d - kShadowEps)) return true;
  float alb[3], eml[3];
  mat_albedo0(a, m, alb);
  mat_emission0(a, a.mat_ids[lface], eml);
  float S = kInvPi * cx * cy / d2 * a.total_light_area;
  for (int c = 0; c < 3; ++c) L[c] += alb[c] * S * eml[c];
  return true;
}

__global__ void k_geo_primary(PtdSceneArgs a, const int* face_verts,
                              const int* edges, int n_edges,
                              const float* edge_cdf, GeoCam gc, int H, int W,
                              int n_samples, unsigned long long seed,
                              const float* grad_image, float* grad_verts,
                              float total_edge_len) {
  const float* cam = gc.c;
  int sid = blockIdx.x * blockDim.x + threadIdx.x;
  if (sid >= n_samples) return;
  curandStatePhilox4_32_10_t st;
  curand_init(seed ^ 0xda3e39cb94b95bdbull, (unsigned long long)sid, 0, &st);
  float r0 = curand_uniform(&st);
  float rs = curand_uniform(&st);
  float r1 = curand_uniform(&st);
  float r2 = curand_uniform(&st);
  float r3 = curand_uniform(&st);
  int ei = cdf_pick_geo(edge_cdf, n_edges, r0);
  const int* E = &edges[ei * 4];
  int fa = E[2];
  int fb = E[3];
  int ca = E[0] & 3;
  int cb = E[1] & 3;
  int va = face_verts[fa * 3 + ca];
  int vb = face_verts[fa * 3 + cb];
  V3 p0 = vert(a, fa, ca);
  V3 p1 = vert(a, fa, cb);
  V3 pe = p0 + (p1 - p0) * rs;

  V3 o = v3(cam[0], cam[1], cam[2]);
  // silhouette from the camera origin
  V3 nga = vcross(vert(a, fa, 1) - vert(a, fa, 0),
                  vert(a, fa, 2) - vert(a, fa, 0));
  float sa_side = vdot(nga, o - vert(a, fa, 0));
  bool sil;
  if (fb < 0) {
    sil = true;
  } else {
    V3 ngb = vcross(vert(a, fb, 1) - vert(a, fb, 0),
                    vert(a, fb, 2) - vert(a, fb, 0));
    float sb_side = vdot(ngb, o - vert(a, fb, 0));
    sil = (sa_side > 0.0f) != (sb_side > 0.0f);
  }
  if (!sil) return;

  // the edge point must be directly visible: nothing strictly in front
  V3 wdir = pe - o;
  float dpe = sqrtf(fmaxf(vdot(wdir, wdir), 1e-12f));
  wdir = wdir * (1.0f / dpe);
  {
    float tb = dpe - fmaxf(kShadowEps, 1e-3f * dpe);
    float bu, bv;
    if (geo_closest(a, o, wdir, kRayEps, tb, bu, bv) >= 0) return;
  }

  // image-space projection Jacobian at pe: q = pe - o, wplane = q/<q,f>
  // (unit forward), J(dq) = ((dq - wplane <f,dq>)/<q,f>) resolved on the
  // rs/us axes and scaled to pixel units.
  V3 fwd = v3(cam[3], cam[4], cam[5]);
  V3 rsv = v3(cam[6], cam[7], cam[8]);
  V3 usv = v3(cam[9], cam[10], cam[11]);
  float rs2 = vdot(rsv, rsv), us2 = vdot(usv, usv);
  V3 q = pe - o;
  float qf = vdot(q, fwd);
  if (qf <= 1e-6f) return;  // behind the camera
  V3 wplane = q * (1.0f / qf);
  float nx = vdot(wplane - fwd, rsv) / rs2;
  float ny = vdot(wplane - fwd, usv) / us2;
  float pxf = (nx + 1.0f) * 0.5f * (float)W;
  float pyf = (1.0f - ny) * 0.5f * (float)H;
  int pxi = (int)pxf, pyi = (int)pyf;
  if (pxf < 0.0f || pxi >= W || pyf < 0.0f || pyi >= H) return;
  const float* gpix = &grad_image[(pyi * W + pxi) * 3];

  auto jac_x = [&](V3 dq) {
    return 0.5f * (float)W * vdot(dq - wplane * vdot(fwd, dq), rsv) /
           (rs2 * qf);
  };
  auto jac_y = [&](V3 dq) {
    return -0.5f * (float)H * vdot(dq - wplane * vdot(fwd, dq), usv) /
           (us2 * qf);
  };
  V3 edir = vnorm(p1 - p0);
  float tx = jac_x(edir), ty = jac_y(edir);
  float tl = sqrtf(tx * tx + ty * ty);  // image arc length per edge arc
  if (tl < 1e-9f) return;
  float n2x = -ty / tl, n2y = tx / tl;  // unit image normal to the curve

  // probe both sides one pixel off: classify by whether the first hit is
  // the edge's own face pair (front) or something else (background)
  bool istop[2];
  for (int sdx = 0; sdx < 2; ++sdx) {
    float sp = (sdx == 0) ? 1.0f : -1.0f;
    float ppx = pxf + sp * n2x;
    float ppy = pyf + sp * n2y;
    float nnx = 2.0f * ppx / (float)W - 1.0f;
    float nny = 1.0f - 2.0f * ppy / (float)H;
    V3 dir = vnorm(v3(fwd.x + nnx * rsv.x + nny * usv.x,
                      fwd.y + nnx * rsv.y + nny * usv.y,
                      fwd.z + nnx * rsv.z + nny * usv.z));
    float th = 1e30f;
    float bu, bv;
    int fh = geo_closest(a, o, dir, kRayEps, th, bu, bv);
    istop[sdx] = (fh == fa) || (fb >= 0 && fh == fb);
  }
  if (istop[0] == istop[1]) return;  // not a clean outline crossing
  float sign = istop[0] ? -1.0f : 1.0f;  // n_img points into the background

  // background: the exact limit point, continuing the camera ray past pe
  int fhit_bg;
  V3 xbg = v3(0, 0, 0);
  {
    float tb2 = 1e30f;
    float bu, bv;
    fhit_bg = geo_closest(a, o, wdir, dpe + fmaxf(kShadowEps, 1e-3f * dpe),
                          tb2, bu, bv);
    if (fhit_bg == fa || (fb >= 0 && fhit_bg == fb)) return;  // grazing
    if (fhit_bg >= 0) xbg = o + wdir * tb2;
  }

  // direct radiance on both sides, same light draw
  float Ltop[3], Lbg[3];
  int ftop = fa;
  if (fb >= 0) {
    // camera-facing member of the pair
    if (sa_side <= 0.0f) ftop = fb;
  }
  if (!geo_direct_radiance(a, ftop, pe, o, r1, r2, r3, Ltop)) return;
  if (!geo_direct_radiance(a, fhit_bg, xbg, o, r1, r2, r3, Lbg)) return;

  float base = 0.0f;
  for (int c = 0; c < 3; ++c) base += gpix[c] * (Ltop[c] - Lbg[c]);
  base *= sign * tl * total_edge_len / (float)n_samples;

  for (int vi = 0; vi < 2; ++vi) {
    float wgt = (vi == 0) ? (1.0f - rs) : rs;
    int target = (vi == 0) ? va : vb;
    for (int cc = 0; cc < 3; ++cc) {
      V3 dq = v3(cc == 0 ? 1.0f : 0.0f, cc == 1 ? 1.0f : 0.0f,
                 cc == 2 ? 1.0f : 0.0f) * wgt;
      float vel = n2x * jac_x(dq) + n2y * jac_y(dq);
      if (vel != 0.0f)
        atomicAdd(&grad_verts[target * 3 + cc], base * vel);
    }
  }
}
}  // namespace

extern "C" void ptd_geo_interior_launch(const PtdSceneArgs* args,
                                        const int* face_verts,
                                        const float* cam, int H, int W,
                                        int spp, long long seed,
                                        const float* grad_image,
                                        float* grad_verts,
                                        cudaStream_t stream) {
  GeoCam gc;
  for (int i = 0; i < 12; ++i) gc.c[i] = cam[i];
  int blocks = (H * W + kThreads - 1) / kThreads;
  k_geo_interior<<<blocks, kThreads, 0, stream>>>(
      *args, face_verts, gc, H, W, spp, (unsigned long long)seed,
      grad_image, grad_verts);
}

extern "C" void ptd_geo_boundary_launch(const PtdSceneArgs* args,
                                        const int* face_verts,
                                        const int* edges, int n_edges,
                                        const float* edge_cdf,
                                        const float* cam, int H, int W,
                                        int n_samples, long long seed,
                                        const float* grad_image,
                                        float* grad_verts,
                                        cudaStream_t stream) {
  // total edge length rides in edge_cdf[n_edges] (one past the cdf)
  float total_len_h = 0.0f;
  cudaMemcpyAsync(&total_len_h, edge_cdf + n_edges, sizeof(float),
                  cudaMemcpyDeviceToHost, stream);
  cudaStreamSynchronize(stream);
  GeoCam gc;
  for (int i = 0; i < 12; ++i) gc.c[i] = cam[i];
  int blocks = (n_samples + kThreads - 1) / kThreads;
  k_geo_boundary<<<blocks, kThreads, 0, stream>>>(
      *args, face_verts, edges, n_edges, edge_cdf, gc, H, W, n_samples,
      (unsigned long long)seed, grad_image, grad_verts, total_len_h);
  k_geo_primary<<<blocks, kThreads, 0, stream>>>(
      *args, face_verts, edges, n_edges, edge_cdf, gc, H, W, n_samples,
      (unsigned long long)seed, grad_image, grad_verts, total_len_h);
}