// pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
//
// Forward: a megakernel unidirectional path tracer. Materials: Lambertian
// diffuse, GGX conductor (VNDF, height-correlated Smith), smooth dielectric,
// rough plastic (diffuse base + GGX coat, F0 = 0.04), and rough dielectric
// (GGX transmission, BSDF-sampled only). Lights: emissive faces with
// per-texel emission textures, and an importance-sampled equirectangular
// environment map. An optional homogeneous participating medium (absorption
// + isotropic scattering) fills the scene, with distances sampled from a
// detached rate frozen at Scene construction. Estimators: MIS (balance
// heuristic) by default, plus legacy NEE-only and BRDF-only modes.
// Per-texel albedo with bilinear filtering; binned-SAH BVH with
// near-child-first traversal. One thread owns one pixel and accumulates its
// spp samples serially, so the image is bitwise deterministic.
//
// Backward: exact path replay. Sampling depends only on geometry, frozen
// material parameters (roughness, ior), the detached environment CDF, the
// detached medium sampling rate, and the counter-based Philox stream --
// never on the differentiable parameters (albedo texels, emission texels,
// environment texels, sigma_a, sigma_s). The replay re-traces identical
// paths with identical draws and no stored path state. Every radiance term
// is trAcc * (prod of per-bounce factors) * S * E where each factor is
// AFFINE in its vertex's albedo texels (Schlick Fresnel is affine in F0)
// and E is a linear emission or environment texel; medium transmittance
// ratios fold into trAcc and the factors, and their sigma-derivatives are
// the closed-form log-derivatives -(D_total + d_term) and
// -(D_total + d_term) + N_scatter/sigma_s. Gradients scatter through the
// exact bilinear-footprint adjoints. Interior terms only here; geometry
// gradients live in the companion geometry kernels.

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

#include <cstdint>

#include "pathtracer_launch.h"

namespace {

constexpr int kThreads = 128;
constexpr int kMaxBounces = 16;
constexpr int kMaxMats = 64;
constexpr int kSharedTexels = 2048;   // shared-staged albedo grad limit
constexpr int kSharedEmiTexels = 512; // shared-staged emission grad limit
constexpr int kStack = 64;
constexpr float kPi = 3.14159265358979323846f;
constexpr float kInvPi = 0.31830988618379067154f;
constexpr float kInv4Pi = 0.07957747154594766788f;
constexpr float kRayEps = 1e-4f;
constexpr float kShadowEps = 1e-3f;
constexpr float kEnvDist = 1e30f;

enum Mode { kModeBrdf = 0, kModeNee = 1, kModeMis = 2 };
enum MatType {
  kDiffuse = 0,
  kConductor = 1,
  kDielectric = 2,       // smooth
  kPlastic = 3,          // diffuse base + GGX coat, F0 = 0.04
  kRoughDielectric = 4,  // GGX transmission, BSDF-sampled only
};
constexpr float kF0Coat = 0.04f;

// ---------------------------------------------------------------- float3 ops
__device__ __forceinline__ float3 f3(float x, float y, float z) {
  return make_float3(x, y, z);
}
__device__ __forceinline__ float3 operator+(float3 a, float3 b) {
  return f3(a.x + b.x, a.y + b.y, a.z + b.z);
}
__device__ __forceinline__ float3 operator-(float3 a, float3 b) {
  return f3(a.x - b.x, a.y - b.y, a.z - b.z);
}
__device__ __forceinline__ float3 operator*(float3 a, float s) {
  return f3(a.x * s, a.y * s, a.z * s);
}
__device__ __forceinline__ float dot(float3 a, float3 b) {
  return a.x * b.x + a.y * b.y + a.z * b.z;
}
__device__ __forceinline__ float3 cross(float3 a, float3 b) {
  return f3(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__ float3 normalize(float3 a) {
  return a * rsqrtf(fmaxf(dot(a, a), 1e-30f));
}
__device__ __forceinline__ float3 neg(float3 a) { return f3(-a.x, -a.y, -a.z); }

// Duff et al. 2017, "Building an Orthonormal Basis, Revisited".
__device__ __forceinline__ void onb(float3 n, float3& t, float3& b) {
  float s = copysignf(1.0f, n.z);
  float a = -1.0f / (s + n.z);
  float c = n.x * n.y * a;
  t = f3(1.0f + s * n.x * n.x * a, s * c, -s * n.x);
  b = f3(c, s + n.y * n.y * a, -n.y);
}

// ---------------------------------------------------------------- scene view
struct DevScene {
  const float* tris;   // [F, 9]
  const int* mats;     // [F]
  const float* uvs;    // [F, 3, 2]
  const float* nf;     // [N, 6]
  const int* ni;       // [N, 3] internal (l, r, axis<<1) | leaf (start, count, odd)
  int n_nodes;
  const int* lf;       // [L] emissive faces
  const float* lcdf;   // [L]
  int nl;
  float larea;
  const float* tex;    // [T, 3] albedo texels
  const int* thdr;     // [M, 3] (offset, W, H)
  int nt;
  const float* etex;   // [Te, 3] emission texels
  const int* ehdr;     // [M, 3] (offset, W, H)
  int net;
  const int* mtype;    // [M]
  const float* mrough; // [M]
  const float* mior;   // [M]
  int nm;
  const float* med_sa; // [3] live absorption, or null
  const float* med_ss; // [3] live scattering
  float med_sbar;      // frozen (detached) sampling rate
  int has_med;
  const float* env;    // [Eh*Ew, 3] or null
  int ew, eh;
  const float* ecdf_m;
  const float* ecdf_c;
  const float* epdf;
};

struct DevCam {
  float p[3], f[3], r[3], u[3];
};

// grad buffers bundled to keep call signatures sane
struct GCtx {
  const float* gs;   // per-pixel dLoss/dpixel / spp
  float* g_tex;      // [T, 3] global
  float* s_gt;       // shared staging for g_tex or null
  float* g_etex;     // [Te, 3] global
  float* s_get;      // shared staging for g_etex or null
  float* g_env;      // [Eh*Ew, 3] global or null
  float* s_gm;       // shared [6]: d sigma_a[3], d sigma_s[3]
};

// ------------------------------------------------------------- intersection
__device__ __forceinline__ bool tri_hit(const float* v, float3 ro, float3 rd,
                                        float tmin, float tmax, float& t,
                                        float& bu, float& bv, float3& ng) {
  float3 v0 = f3(v[0], v[1], v[2]);
  float3 e1 = f3(v[3], v[4], v[5]) - v0;
  float3 e2 = f3(v[6], v[7], v[8]) - v0;
  float3 p = cross(rd, e2);
  float det = dot(e1, p);
  if (fabsf(det) < 1e-12f) return false;
  float inv = 1.0f / det;
  float3 s = ro - v0;
  float u = dot(s, p) * inv;
  if (u < -1e-6f || u > 1.0f + 1e-6f) return false;
  float3 q = cross(s, e1);
  float w = dot(rd, q) * inv;
  if (w < -1e-6f || u + w > 1.0f + 1e-6f) return false;
  float tt = dot(e2, q) * inv;
  if (tt < tmin || tt > tmax) return false;
  t = tt;
  bu = u;
  bv = w;
  ng = cross(e1, e2);
  return true;
}

__device__ __forceinline__ bool slab(const float* b, const float ro[3],
                                     const float inv[3], float tmax) {
  float t0 = kRayEps, t1 = tmax;
#pragma unroll
  for (int a = 0; a < 3; ++a) {
    float lo = (b[a] - ro[a]) * inv[a];
    float hi = (b[3 + a] - ro[a]) * inv[a];
    if (lo > hi) {
      float tmp = lo;
      lo = hi;
      hi = tmp;
    }
    t0 = fmaxf(t0, lo);
    t1 = fminf(t1, hi);
  }
  return t0 <= t1;
}

__device__ __forceinline__ void inv_dir(float3 rd, 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;
}

__device__ int bvh_closest(const DevScene& sc, float3 ro, float3 rd,
                           float tmin, float& tbest, float& bu, float& bv,
                           float3& ngbest) {
  float roa[3] = {ro.x, ro.y, ro.z};
  float dira[3] = {rd.x, rd.y, rd.z};
  float inv[3];
  inv_dir(rd, inv);
  int stack[kStack];
  int sp = 0;
  stack[sp++] = 0;
  int best = -1;
  while (sp > 0) {
    int nid = stack[--sp];
    if (!slab(&sc.nf[nid * 6], roa, inv, tbest)) continue;
    const int* n = &sc.ni[nid * 3];
    if (n[2] & 1) {
      for (int f = n[0]; f < n[0] + n[1]; ++f) {
        float t, u, v;
        float3 ng;
        if (tri_hit(&sc.tris[f * 9], ro, rd, tmin, tbest, t, u, v, ng)) {
          tbest = t;
          bu = u;
          bv = v;
          ngbest = ng;
          best = f;
        }
      }
    } else if (sp + 2 <= kStack) {
      int axis = n[2] >> 1;
      int near = (dira[axis] >= 0.0f) ? n[0] : n[1];
      int far = (dira[axis] >= 0.0f) ? n[1] : n[0];
      stack[sp++] = far;
      stack[sp++] = near;
    }
  }
  return best;
}

__device__ bool bvh_occluded(const DevScene& sc, float3 ro, float3 rd,
                             float tmax) {
  float roa[3] = {ro.x, ro.y, ro.z};
  float inv[3];
  inv_dir(rd, inv);
  int stack[kStack];
  int sp = 0;
  stack[sp++] = 0;
  while (sp > 0) {
    int nid = stack[--sp];
    if (!slab(&sc.nf[nid * 6], roa, inv, tmax)) continue;
    const int* n = &sc.ni[nid * 3];
    if (n[2] & 1) {
      for (int f = n[0]; f < n[0] + n[1]; ++f) {
        float t, u, v;
        float3 ng;
        if (tri_hit(&sc.tris[f * 9], ro, rd, kRayEps, tmax, t, u, v, ng))
          return true;
      }
    } else if (sp + 2 <= kStack) {
      stack[sp++] = n[0];
      stack[sp++] = n[1];
    }
  }
  return false;
}

__device__ __forceinline__ int cdf_pick(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;
}

// Bilinear fetch with repeat wrap over a flat texel block.
__device__ __forceinline__ void bilinear(const float* block, int off, int W,
                                         int H, float u, float v, float rgb[3],
                                         int idx4[4], float w4[4]) {
  float uu = u * (float)W - 0.5f;
  float vv = v * (float)H - 0.5f;
  float fx = uu - floorf(uu);
  float fy = vv - floorf(vv);
  int x0 = (int)floorf(uu), y0 = (int)floorf(vv);
  int x0w = ((x0 % W) + W) % W;
  int x1w = (((x0 + 1) % W) + W) % W;
  int y0w = ((y0 % H) + H) % H;
  int y1w = (((y0 + 1) % H) + H) % H;
  idx4[0] = off + y0w * W + x0w; w4[0] = (1.0f - fx) * (1.0f - fy);
  idx4[1] = off + y0w * W + x1w; w4[1] = fx * (1.0f - fy);
  idx4[2] = off + y1w * W + x0w; w4[2] = (1.0f - fx) * fy;
  idx4[3] = off + y1w * W + x1w; w4[3] = fx * fy;
  rgb[0] = rgb[1] = rgb[2] = 0.0f;
#pragma unroll
  for (int i = 0; i < 4; ++i) {
    const float* t = &block[idx4[i] * 3];
    rgb[0] += w4[i] * t[0];
    rgb[1] += w4[i] * t[1];
    rgb[2] += w4[i] * t[2];
  }
}

__device__ __forceinline__ void face_uv(const DevScene& sc, int face,
                                        float bu, float bv, float& u,
                                        float& v) {
  const float* U = &sc.uvs[face * 6];
  float w0 = 1.0f - bu - bv;
  u = w0 * U[0] + bu * U[2] + bv * U[4];
  v = w0 * U[1] + bu * U[3] + bv * U[5];
}

__device__ __forceinline__ void sample_albedo(const DevScene& sc, int m,
                                              int face, float bu, float bv,
                                              float rgb[3], int idx4[4],
                                              float w4[4]) {
  float u, v;
  face_uv(sc, face, bu, bv, u, v);
  const int* h = &sc.thdr[m * 3];
  bilinear(sc.tex, h[0], h[1], h[2], u, v, rgb, idx4, w4);
}

__device__ __forceinline__ void sample_emission(const DevScene& sc, int m,
                                                int face, float bu, float bv,
                                                float rgb[3], int idx4[4],
                                                float w4[4]) {
  float u, v;
  face_uv(sc, face, bu, bv, u, v);
  const int* h = &sc.ehdr[m * 3];
  bilinear(sc.etex, h[0], h[1], h[2], u, v, rgb, idx4, w4);
}

// ----------------------------------------------------------- environment map
__device__ __forceinline__ void dir_to_equirect(float3 d, float& u, float& v) {
  float phi = atan2f(d.z, d.x);
  float theta = acosf(fminf(fmaxf(d.y, -1.0f), 1.0f));
  u = (phi + kPi) / (2.0f * kPi);
  v = theta / kPi;
}

__device__ __forceinline__ float3 equirect_to_dir(float u, float v) {
  float phi = u * 2.0f * kPi - kPi;
  float theta = v * kPi;
  float st = sinf(theta);
  return f3(st * cosf(phi), cosf(theta), st * sinf(phi));
}

__device__ __forceinline__ float env_pdf(const DevScene& sc, float3 d) {
  float u, v;
  dir_to_equirect(d, u, v);
  int x = min(sc.ew - 1, (int)(u * sc.ew));
  int y = min(sc.eh - 1, (int)(v * sc.eh));
  float st = fmaxf(sinf((y + 0.5f) * kPi / sc.eh), 1e-4f);
  float p_img = sc.epdf[y * sc.ew + x];
  float p_tab = p_img * (float)(sc.ew * sc.eh) / (2.0f * kPi * kPi * st);
  return 0.5f * p_tab + 0.5f * (1.0f / (4.0f * kPi));
}

__device__ __forceinline__ void env_fetch(const DevScene& sc, float3 d,
                                          float rgb[3], int idx4[4],
                                          float w4[4]) {
  float u, v;
  dir_to_equirect(d, u, v);
  v = fminf(fmaxf(v, 0.5f / sc.eh), 1.0f - 0.5f / sc.eh);
  bilinear(sc.env, 0, sc.ew, sc.eh, u, v, rgb, idx4, w4);
}

// ------------------------------------------------------------------- GGX
__device__ __forceinline__ float ggx_lambda(float a2, float cs) {
  cs = fabsf(cs);
  float c2 = cs * cs;
  float t2 = fmaxf(0.0f, 1.0f - c2) / fmaxf(c2, 1e-12f);
  return 0.5f * (-1.0f + sqrtf(1.0f + a2 * t2));
}

__device__ __forceinline__ float ggx_d(float a2, float ch) {
  float d = ch * ch * (a2 - 1.0f) + 1.0f;
  return a2 / fmaxf(kPi * d * d, 1e-20f);
}

__device__ __forceinline__ float3 ggx_sample_vndf(float3 wi, float alpha,
                                                  float u1, float u2) {
  float3 vh = normalize(f3(alpha * wi.x, alpha * wi.y, wi.z));
  float lensq = vh.x * vh.x + vh.y * vh.y;
  float3 T1 = lensq > 1e-12f ? f3(-vh.y, vh.x, 0.0f) * rsqrtf(lensq)
                             : f3(1.0f, 0.0f, 0.0f);
  float3 T2 = cross(vh, T1);
  float r = sqrtf(u1);
  float phi = 2.0f * kPi * u2;
  float t1 = r * cosf(phi);
  float t2 = r * sinf(phi);
  float s = 0.5f * (1.0f + vh.z);
  t2 = (1.0f - s) * sqrtf(fmaxf(0.0f, 1.0f - t1 * t1)) + s * t2;
  float3 nh = T1 * t1 + T2 * t2 +
              vh * sqrtf(fmaxf(0.0f, 1.0f - t1 * t1 - t2 * t2));
  return normalize(f3(alpha * nh.x, alpha * nh.y, fmaxf(1e-6f, nh.z)));
}

__device__ __forceinline__ float ggx_pdf(float a2, float3 wi, float3 wo) {
  float3 h = normalize(wi + wo);
  float ch = fmaxf(h.z, 1e-6f);
  float wih = fmaxf(dot(wi, h), 1e-6f);
  float g1 = 1.0f / (1.0f + ggx_lambda(a2, wi.z));
  return g1 * ggx_d(a2, ch) * wih / fmaxf(wi.z, 1e-6f) / (4.0f * wih);
}

// f_spec * cos_o for a GGX lobe with Fresnel factor Fc (already evaluated):
// D * Fc * G2 / (4 wi.z). The caller divides by pdf or folds geometry.
__device__ __forceinline__ float ggx_spec_cos(float a2, float3 wi, float3 wo,
                                              float& sgl_out) {
  float3 h = normalize(wi + wo);
  sgl_out = 0.0f;
  float m = fminf(fmaxf(1.0f - fmaxf(dot(wi, h), 0.0f), 0.0f), 1.0f);
  float m2 = m * m;
  sgl_out = m2 * m2 * m;
  float G2 = 1.0f / (1.0f + ggx_lambda(a2, wi.z) + ggx_lambda(a2, wo.z));
  return ggx_d(a2, fmaxf(h.z, 1e-6f)) * G2 / (4.0f * fmaxf(wi.z, 1e-6f));
}

__device__ __forceinline__ float fresnel_dielectric(float cos_i, float eta) {
  // eta = n2/n1
  float s2 = (1.0f - cos_i * cos_i) / (eta * eta);
  if (s2 >= 1.0f) return 1.0f;
  float cos_t = sqrtf(1.0f - s2);
  float rs = (cos_i - eta * cos_t) / (cos_i + eta * cos_t);
  float rp = (eta * cos_i - cos_t) / (eta * cos_i + cos_t);
  return 0.5f * (rs * rs + rp * rp);
}

__device__ __forceinline__ float schlick_s(float cos_h) {
  float m = fminf(fmaxf(1.0f - cos_h, 0.0f), 1.0f);
  float m2 = m * m;
  return m2 * m2 * m;
}

__device__ __forceinline__ float mis_w(float pa, float pb) {
  return pa / fmaxf(pa + pb, 1e-20f);
}

__device__ __forceinline__ bool nee_capable(int mt) {
  return mt == kDiffuse || mt == kConductor || mt == kPlastic;
}

// ------------------------------------------------------------- path tracing
__device__ __forceinline__ void scatter_buf(float* shared_buf, float* global_buf,
                                            int idx, int c, float v) {
  if (shared_buf) atomicAdd(&shared_buf[idx * 3 + c], v);
  else atomicAdd(&global_buf[idx * 3 + c], v);
}

struct Factor {
  float val[3];
  float dc[3];
  int fi[4];
  float fw[4];
};

// One radiance term: trAcc * (prod of nk factors, tf last when non-null) *
// S * E, with E from an emission footprint (e_env=false), an env footprint
// (e_env=true), or plain (ei null). dsa/dss are the medium log-derivative
// sums for this term (0 when no medium).
template <bool GRAD>
__device__ void add_term(const Factor* fs, int nk, const float T[3],
                         const float* tf, float S, const float E[3],
                         const int* ei, const float* ew, bool e_env,
                         const float trAcc[3], float dsa, const float dss[3],
                         float3& acc, const GCtx& g) {
  float P[3] = {T[0], T[1], T[2]};
  if (tf) {
    P[0] *= tf[0];
    P[1] *= tf[1];
    P[2] *= tf[2];
  }
  if (!GRAD) {
    acc.x += P[0] * S * E[0];
    acc.y += P[1] * S * E[1];
    acc.z += P[2] * S * E[2];
    return;
  }
  // E-side gradients
  if (ei) {
    for (int c = 0; c < 3; ++c) {
      float base = g.gs[c] * P[c] * S;
#pragma unroll
      for (int i = 0; i < 4; ++i) {
        if (e_env) atomicAdd(&g.g_env[ei[i] * 3 + c], base * ew[i]);
        else scatter_buf(g.s_get, g.g_etex, ei[i], c, base * ew[i]);
      }
    }
  }
  // medium sigma gradients: g * term * dln
  if (g.s_gm && (dsa != 0.0f || dss[0] != 0.0f || dss[1] != 0.0f ||
                 dss[2] != 0.0f)) {
    for (int c = 0; c < 3; ++c) {
      float term = g.gs[c] * P[c] * S * E[c];
      atomicAdd(&g.s_gm[c], term * dsa);
      atomicAdd(&g.s_gm[3 + c], term * dss[c]);
    }
  }
  // albedo texel gradients via exclusion products, scaled by trAcc
  if (nk <= 0) return;
  float suf[kMaxBounces + 1][3];
  suf[nk][0] = suf[nk][1] = suf[nk][2] = 1.0f;
  for (int i = nk - 1; i >= 0; --i)
    for (int c = 0; c < 3; ++c) suf[i][c] = suf[i + 1][c] * fs[i].val[c];
  float pref[3] = {1.0f, 1.0f, 1.0f};
  for (int j = 0; j < nk; ++j) {
    for (int c = 0; c < 3; ++c) {
      if (fs[j].dc[c] != 0.0f) {
        float base = g.gs[c] * E[c] * S * trAcc[c] * pref[c] *
                     suf[j + 1][c] * fs[j].dc[c];
#pragma unroll
        for (int i = 0; i < 4; ++i)
          scatter_buf(g.s_gt, g.g_tex, fs[j].fi[i], c, base * fs[j].fw[i]);
      }
      pref[c] *= fs[j].val[c];
    }
  }
}

template <bool GRAD>
__device__ void trace_path(const DevScene& sc, float3 ro, float3 rd,
                           curandStatePhilox4_32_10_t& st, int B, int mode,
                           float3& acc, const GCtx& g) {
  float T[3] = {1.0f, 1.0f, 1.0f};   // includes bsdf factors AND trAcc
  float trAcc[3] = {1.0f, 1.0f, 1.0f};
  Factor fs[kMaxBounces];
  int nk = 0;
  float prev_pdf = 0.0f;
  bool prev_delta = true;
  bool has_env = sc.env != nullptr;
  float Dtot = 0.0f;   // medium distance so far
  int Nsc = 0;         // medium scatter events so far
  float sa[3] = {0, 0, 0}, ss[3] = {0, 0, 0}, stt[3] = {0, 0, 0};
  if (sc.has_med) {
    for (int c = 0; c < 3; ++c) {
      sa[c] = sc.med_sa[c];
      ss[c] = fmaxf(sc.med_ss[c], 1e-8f);
      stt[c] = sa[c] + ss[c];
    }
  }

  auto med_dsa = [&](float dterm) { return sc.has_med ? -(Dtot + dterm) : 0.0f; };

  for (int k = 0; k < B; ++k) {
    float dmed = kEnvDist;
    if (sc.has_med) {
      float u = curand_uniform(&st);  // in (0, 1]
      dmed = -logf(fmaxf(u, 1e-12f)) / sc.med_sbar;
    }

    float tbest = kEnvDist;
    float bu = 0.0f, bv = 0.0f;
    float3 ng;
    int face = bvh_closest(sc, ro, rd, kRayEps, tbest, bu, bv, ng);

    // ----------------------------------------------------- medium vertex
    if (sc.has_med && dmed < tbest) {
      Dtot += dmed;
      float3 x = ro + rd * dmed;
      // event factor: sigma_s * exp(-sigma_t d) / (sbar * exp(-sbar d))
      float denom = sc.med_sbar * expf(-sc.med_sbar * dmed);
      if (nk >= kMaxBounces) break;
      Factor& f = fs[nk++];
      for (int c = 0; c < 3; ++c) {
        f.val[c] = ss[c] * expf(-stt[c] * dmed) / denom;
        f.dc[c] = 0.0f;
        T[c] *= f.val[c];
      }
#pragma unroll
      for (int i = 0; i < 4; ++i) { f.fi[i] = 0; f.fw[i] = 0.0f; }
      Nsc += 1;

      // light NEE from the medium point (isotropic phase, no cosine at x)
      bool draws_light = (mode != kModeBrdf) && sc.nl > 0 &&
                         (mode == kModeMis || k + 1 < B);
      if (draws_light) {
        float r1 = curand_uniform(&st);
        float r2 = curand_uniform(&st);
        float r3 = curand_uniform(&st);
        int li = cdf_pick(sc.lcdf, sc.nl, r1);
        int lface = sc.lf[li];
        const float* lv = &sc.tris[lface * 9];
        float3 lv0 = f3(lv[0], lv[1], lv[2]);
        float3 le1 = f3(lv[3], lv[4], lv[5]) - lv0;
        float3 le2 = f3(lv[6], lv[7], lv[8]) - lv0;
        float su = sqrtf(r2);
        float b0 = 1.0f - su, b1 = r3 * su;
        float3 y = lv0 + le1 * b0 + le2 * b1;
        float3 ln = normalize(cross(le1, le2));
        float3 dvec = y - x;
        float d2 = fmaxf(dot(dvec, dvec), 1e-12f);
        float d = sqrtf(d2);
        float3 wo_w = dvec * (1.0f / d);
        if (dot(ln, dvec) > 0.0f) ln = neg(ln);
        float cy = -dot(ln, wo_w);
        if (cy > 1e-6f && d > kShadowEps * 2.0f &&
            !bvh_occluded(sc, x, wo_w, d - kShadowEps)) {
          int lm = sc.mats[lface];
          float E[3];
          int ei[4];
          float ewt[4];
          sample_emission(sc, lm, lface, b0, b1, E, ei, ewt);
          float S = kInv4Pi * cy / d2 * sc.larea;
          float trL[3];
          float trm = 1.0f;
          for (int c = 0; c < 3; ++c) trL[c] = expf(-stt[c] * d);
          (void)trm;
          float w = 1.0f;
          if (mode == kModeMis) {
            float p_lw = d2 / fmaxf(cy * sc.larea, 1e-12f);
            w = mis_w(p_lw, kInv4Pi);
          }
          // fold shadow transmittance into the term factor
          float tf[3] = {trL[0], trL[1], trL[2]};
          float dss_t[3];
          float dsa_t = med_dsa(d);
          for (int c = 0; c < 3; ++c) dss_t[c] = dsa_t + Nsc / ss[c];
          add_term<GRAD>(fs, nk, T, tf, S * w, E, ei, ewt, false, trAcc,
                         dsa_t, dss_t, acc, g);
        }
      }
      // isotropic phase continuation
      float p1 = curand_uniform(&st);
      float p2 = curand_uniform(&st);
      float z = 1.0f - 2.0f * p1;
      float rxy = sqrtf(fmaxf(0.0f, 1.0f - z * z));
      float ph = 2.0f * kPi * p2;
      rd = f3(rxy * cosf(ph), z, rxy * sinf(ph));
      ro = x;
      prev_delta = false;
      prev_pdf = kInv4Pi;
      continue;
    }

    // --------------------------------------------------- surface / miss
    if (sc.has_med && face >= 0) {
      // survival ratio to the surface: exp(-sigma_t t) / exp(-sbar t)
      Dtot += tbest;
      float esb = expf(sc.med_sbar * tbest);  // 1 / exp(-sbar t)
      for (int c = 0; c < 3; ++c) {
        float r = expf(-stt[c] * tbest) * esb;
        trAcc[c] *= r;
        T[c] *= r;
      }
    }

    if (face < 0) {
      if (has_env && !sc.has_med) {
        float E[3];
        int ei[4];
        float ewd[4];
        env_fetch(sc, rd, E, ei, ewd);
        float w = 1.0f;
        if (mode == kModeMis && !prev_delta)
          w = mis_w(prev_pdf, env_pdf(sc, rd));
        if (mode != kModeNee || prev_delta) {
          float z3[3] = {0, 0, 0};
          add_term<GRAD>(fs, nk, T, nullptr, w, E, ei, ewd, true, trAcc,
                         0.0f, z3, acc, g);
        }
      }
      break;
    }

    int m = sc.mats[face];
    int mt = sc.mtype[m];
    float3 x = ro + rd * tbest;
    float3 n = normalize(ng);
    bool backface = dot(n, rd) > 0.0f;
    if (backface) n = neg(n);

    float am[3];
    int idx4[4];
    float w4[4];
    sample_albedo(sc, m, face, bu, bv, am, idx4, w4);

    // emission gather
    {
      float E[3];
      int ei[4];
      float ewt[4];
      sample_emission(sc, m, face, bu, bv, E, ei, ewt);
      if (E[0] > 0.0f || E[1] > 0.0f || E[2] > 0.0f) {
        float w = 1.0f;
        bool add = true;
        if (mode == kModeNee) {
          add = (k == 0);
        } else if (mode == kModeMis && !prev_delta && sc.nl > 0) {
          float cy = fabsf(dot(normalize(ng), rd));
          float p_l = (tbest * tbest) / fmaxf(cy * sc.larea, 1e-12f);
          w = mis_w(prev_pdf, p_l);
        }
        if (add) {
          float dss_t[3];
          float dsa_t = med_dsa(0.0f);
          for (int c = 0; c < 3; ++c)
            dss_t[c] = sc.has_med ? dsa_t + Nsc / ss[c] : 0.0f;
          add_term<GRAD>(fs, nk, T, nullptr, w, E, ei, ewt, false, trAcc,
                         dsa_t, dss_t, acc, g);
        }
      }
    }

    // ---------------- smooth dielectric: delta lobes
    if (mt == kDielectric) {
      float c1 = curand_uniform(&st);
      float c2 = curand_uniform(&st);
      float c3 = curand_uniform(&st);
      (void)c2;
      (void)c3;
      float eta_r = backface ? sc.mior[m] : 1.0f / sc.mior[m];  // n1/n2
      float ci = -dot(rd, n);
      float F = fresnel_dielectric(ci, 1.0f / eta_r);
      float3 nd;
      if (c1 < F) {
        nd = rd + n * (2.0f * ci);
        ro = x + n * kRayEps;
      } else {
        float s2 = eta_r * eta_r * (1.0f - ci * ci);
        float ct = sqrtf(fmaxf(0.0f, 1.0f - s2));
        nd = rd * eta_r + n * (eta_r * ci - ct);
        ro = x - n * kRayEps;
      }
      rd = normalize(nd);
      if (nk < kMaxBounces) {
        Factor& f = fs[nk++];
        for (int c = 0; c < 3; ++c) { f.val[c] = 1.0f; f.dc[c] = 0.0f; }
#pragma unroll
        for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
      }
      prev_delta = true;
      prev_pdf = 0.0f;
      continue;
    }

    // local frame
    float3 tang, bit;
    onb(n, tang, bit);
    float3 wi_w = neg(rd);
    float3 wi = f3(dot(wi_w, tang), dot(wi_w, bit), dot(wi_w, n));
    wi.z = fmaxf(wi.z, 1e-6f);
    float alpha = fmaxf(sc.mrough[m], 0.01f);
    float a2 = alpha * alpha;

    // ---------------- rough dielectric: BSDF-sampled only
    if (mt == kRoughDielectric) {
      float c1 = curand_uniform(&st);
      float c2 = curand_uniform(&st);
      float c3 = curand_uniform(&st);
      float3 h = ggx_sample_vndf(wi, alpha, c1, c2);
      float cih = dot(wi, h);
      float eta_r = backface ? sc.mior[m] : 1.0f / sc.mior[m];  // n1/n2
      float F = fresnel_dielectric(fmaxf(cih, 1e-6f), 1.0f / eta_r);
      float3 wo;
      bool transmit = false;
      if (c3 < F) {
        wo = h * (2.0f * cih) - wi;
        if (wo.z <= 1e-6f) break;
      } else {
        float s2 = eta_r * eta_r * (1.0f - cih * cih);
        if (s2 >= 1.0f) {  // numerical TIR guard
          wo = h * (2.0f * cih) - wi;
          if (wo.z <= 1e-6f) break;
        } else {
          float ct = sqrtf(1.0f - s2);
          wo = neg(wi) * eta_r + h * (eta_r * cih - ct);
          if (wo.z >= -1e-6f) break;
          transmit = true;
        }
      }
      float li = ggx_lambda(a2, wi.z);
      float lo = ggx_lambda(a2, wo.z);
      float gw = (1.0f + li) / (1.0f + li + lo);  // G2/G1
      if (nk < kMaxBounces) {
        Factor& f = fs[nk++];
        for (int c = 0; c < 3; ++c) { f.val[c] = gw; f.dc[c] = 0.0f; }
#pragma unroll
        for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
        T[0] *= gw;
        T[1] *= gw;
        T[2] *= gw;
      }
      rd = normalize(tang * wo.x + bit * wo.y + n * wo.z);
      ro = transmit ? x - n * kRayEps : x + n * kRayEps;
      prev_delta = true;  // no NEE partner: gathers take w = 1
      prev_pdf = 0.0f;
      continue;
    }

    // ---------------- light NEE (diffuse / conductor / plastic)
    if ((mode != kModeBrdf) && sc.nl > 0 &&
        (mode == kModeMis || k + 1 < B)) {
      float r1 = curand_uniform(&st);
      float r2 = curand_uniform(&st);
      float r3 = curand_uniform(&st);
      int li = cdf_pick(sc.lcdf, sc.nl, r1);
      int lface = sc.lf[li];
      const float* lv = &sc.tris[lface * 9];
      float3 lv0 = f3(lv[0], lv[1], lv[2]);
      float3 le1 = f3(lv[3], lv[4], lv[5]) - lv0;
      float3 le2 = f3(lv[6], lv[7], lv[8]) - lv0;
      float su = sqrtf(r2);
      float b0 = 1.0f - su, b1 = r3 * su;
      float3 y = lv0 + le1 * b0 + le2 * b1;
      float3 ln = normalize(cross(le1, le2));
      float3 dvec = y - x;
      float d2 = fmaxf(dot(dvec, dvec), 1e-12f);
      float d = sqrtf(d2);
      float3 wo_w = dvec * (1.0f / d);
      if (dot(ln, dvec) > 0.0f) ln = neg(ln);
      float cx = dot(n, wo_w);
      float cy = -dot(ln, wo_w);
      if (cx > 1e-6f && cy > 1e-6f && d > kShadowEps * 2.0f &&
          !bvh_occluded(sc, x + n * kRayEps, wo_w, d - kShadowEps)) {
        int lm = sc.mats[lface];
        float3 wo = f3(dot(wo_w, tang), dot(wo_w, bit), cx);
        float S_geo = cx * cy / d2 * sc.larea;
        float tf_val[3], tf_dc[3];
        float S = S_geo;
        float p_bw = 0.0f;
        if (mt == kDiffuse) {
          S = S_geo * kInvPi;
          for (int c = 0; c < 3; ++c) { tf_val[c] = am[c]; tf_dc[c] = 1.0f; }
          p_bw = fmaxf(wo.z, 0.0f) * kInvPi;
        } else if (mt == kConductor) {
          float sgl;
          float gg = ggx_spec_cos(a2, wi, wo, sgl) / fmaxf(wo.z, 1e-6f);
          // gg here = D G2 / (4 wi.z wo.z); contribution f*cos = gg*wo.z
          S = S_geo * gg;
          for (int c = 0; c < 3; ++c) {
            tf_val[c] = am[c] * (1.0f - sgl) + sgl;
            tf_dc[c] = 1.0f - sgl;
          }
          p_bw = ggx_pdf(a2, wi, wo);
        } else {  // plastic
          float sgl;
          float speco = ggx_spec_cos(a2, wi, wo, sgl);  // D G2 / (4 wi.z)
          float Fc = kF0Coat + (1.0f - kF0Coat) * sgl;
          float spec_f = Fc * speco / fmaxf(wo.z, 1e-6f);
          for (int c = 0; c < 3; ++c) {
            tf_val[c] = am[c] * kInvPi + spec_f;
            tf_dc[c] = kInvPi;
          }
          p_bw = 0.5f * fmaxf(wo.z, 0.0f) * kInvPi + 0.5f * ggx_pdf(a2, wi, wo);
        }
        float w = 1.0f;
        if (mode == kModeMis) {
          float p_lw = d2 / fmaxf(cy * sc.larea, 1e-12f);
          w = mis_w(p_lw, p_bw);
        }
        float E[3];
        int ei[4];
        float ewt[4];
        sample_emission(sc, lm, lface, b0, b1, E, ei, ewt);
        float dsa_t = med_dsa(d);
        float dss_t[3] = {0, 0, 0};
        if (sc.has_med) {
          for (int c = 0; c < 3; ++c) {
            tf_val[c] *= expf(-stt[c] * d);
            tf_dc[c] *= expf(-stt[c] * d);
            dss_t[c] = dsa_t + Nsc / ss[c];
          }
        }
        if (nk < kMaxBounces) {
          Factor& f = fs[nk];
          for (int c = 0; c < 3; ++c) { f.val[c] = tf_val[c]; f.dc[c] = tf_dc[c]; }
#pragma unroll
          for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
          add_term<GRAD>(fs, nk + 1, T, f.val, S * w, E, ei, ewt, false,
                         trAcc, dsa_t, dss_t, acc, g);
        }
      }
    }

    // ---------------- environment NEE (MIS mode only)
    if (mode == kModeMis && has_env) {
      float e1 = curand_uniform(&st);
      float e2 = curand_uniform(&st);
      float e3 = curand_uniform(&st);
      float3 wo_w;
      if (e1 < 0.5f) {
        int y = cdf_pick(sc.ecdf_m, sc.eh, e2);
        int xcol = cdf_pick(&sc.ecdf_c[y * sc.ew], sc.ew, e3);
        float u = (xcol + 0.5f) / sc.ew;
        float v = (y + 0.5f) / sc.eh;
        wo_w = equirect_to_dir(u, v);
      } else {
        float z = 1.0f - 2.0f * e2;
        float rxy = sqrtf(fmaxf(0.0f, 1.0f - z * z));
        float ph = 2.0f * kPi * e3;
        wo_w = f3(rxy * cosf(ph), z, rxy * sinf(ph));
      }
      float cx = dot(n, wo_w);
      if (cx > 1e-6f &&
          !bvh_occluded(sc, x + n * kRayEps, wo_w, kEnvDist)) {
        float p = env_pdf(sc, wo_w);
        float3 wo = f3(dot(wo_w, tang), dot(wo_w, bit), cx);
        float tf_val[3], tf_dc[3];
        float S;
        float p_bw;
        if (mt == kDiffuse) {
          S = cx * kInvPi / p;
          for (int c = 0; c < 3; ++c) { tf_val[c] = am[c]; tf_dc[c] = 1.0f; }
          p_bw = fmaxf(wo.z, 0.0f) * kInvPi;
        } else if (mt == kConductor) {
          float sgl;
          float speco = ggx_spec_cos(a2, wi, wo, sgl);  // D G2 / (4 wi.z)
          S = speco / p;
          for (int c = 0; c < 3; ++c) {
            tf_val[c] = am[c] * (1.0f - sgl) + sgl;
            tf_dc[c] = 1.0f - sgl;
          }
          p_bw = ggx_pdf(a2, wi, wo);
        } else {  // plastic
          float sgl;
          float speco = ggx_spec_cos(a2, wi, wo, sgl);
          float Fc = kF0Coat + (1.0f - kF0Coat) * sgl;
          S = cx / p;
          float spec_f = Fc * speco / fmaxf(cx, 1e-6f);
          for (int c = 0; c < 3; ++c) {
            tf_val[c] = am[c] * kInvPi + spec_f;
            tf_dc[c] = kInvPi;
          }
          p_bw = 0.5f * fmaxf(wo.z, 0.0f) * kInvPi + 0.5f * ggx_pdf(a2, wi, wo);
        }
        float w = mis_w(p, p_bw);
        float E[3];
        int ei[4];
        float ewt[4];
        env_fetch(sc, wo_w, E, ei, ewt);
        float z3[3] = {0, 0, 0};
        if (nk < kMaxBounces) {
          Factor& f = fs[nk];
          for (int c = 0; c < 3; ++c) { f.val[c] = tf_val[c]; f.dc[c] = tf_dc[c]; }
#pragma unroll
          for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
          add_term<GRAD>(fs, nk + 1, T, f.val, S * w, E, ei, ewt, true,
                         trAcc, 0.0f, z3, acc, g);
        }
      }
    }

    // ---------------- BSDF continuation
    float c1 = curand_uniform(&st);
    float c2 = curand_uniform(&st);
    float c3 = curand_uniform(&st);
    float3 wo;
    float fv[3], fdc[3];
    if (mt == kDiffuse) {
      float rr = sqrtf(c1);
      float phi = 2.0f * kPi * c2;
      wo = f3(rr * cosf(phi), rr * sinf(phi), sqrtf(fmaxf(0.0f, 1.0f - c1)));
      prev_pdf = fmaxf(wo.z, 1e-8f) * kInvPi;
      for (int c = 0; c < 3; ++c) { fv[c] = am[c]; fdc[c] = 1.0f; }
    } else if (mt == kConductor) {
      float3 h = ggx_sample_vndf(wi, alpha, c1, c2);
      float wih = dot(wi, h);
      wo = h * (2.0f * wih) - wi;
      if (wo.z <= 1e-6f) break;
      float sgl = schlick_s(fmaxf(wih, 0.0f));
      float li = ggx_lambda(a2, wi.z);
      float lo = ggx_lambda(a2, wo.z);
      float gw = (1.0f + li) / (1.0f + li + lo);
      prev_pdf = ggx_pdf(a2, wi, wo);
      for (int c = 0; c < 3; ++c) {
        fv[c] = (am[c] * (1.0f - sgl) + sgl) * gw;
        fdc[c] = (1.0f - sgl) * gw;
      }
    } else {  // plastic: 50/50 lobe mixture
      if (c3 < 0.5f) {
        float rr = sqrtf(c1);
        float phi = 2.0f * kPi * c2;
        wo = f3(rr * cosf(phi), rr * sinf(phi),
                sqrtf(fmaxf(0.0f, 1.0f - c1)));
      } else {
        float3 h = ggx_sample_vndf(wi, alpha, c1, c2);
        wo = h * (2.0f * dot(wi, h)) - wi;
      }
      if (wo.z <= 1e-6f) break;
      float pdf = 0.5f * wo.z * kInvPi + 0.5f * ggx_pdf(a2, wi, wo);
      prev_pdf = fmaxf(pdf, 1e-10f);
      float sgl;
      float speco = ggx_spec_cos(a2, wi, wo, sgl);
      float Fc = kF0Coat + (1.0f - kF0Coat) * sgl;
      float spec_f = Fc * speco / fmaxf(wo.z, 1e-6f);
      float scale = wo.z / prev_pdf;
      for (int c = 0; c < 3; ++c) {
        fv[c] = (am[c] * kInvPi + spec_f) * scale;
        fdc[c] = kInvPi * scale;
      }
    }
    prev_delta = false;
    rd = normalize(tang * wo.x + bit * wo.y + n * wo.z);
    ro = x + n * kRayEps;
    if (nk < kMaxBounces) {
      Factor& f = fs[nk++];
      for (int c = 0; c < 3; ++c) { f.val[c] = fv[c]; f.dc[c] = fdc[c]; }
#pragma unroll
      for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
    }
    T[0] *= fv[0];
    T[1] *= fv[1];
    T[2] *= fv[2];
  }
}

__device__ __forceinline__ float3 camera_ray(const DevCam& cam, int px, int py,
                                             int W, int H, float jx, float jy) {
  float nx = 2.0f * ((px + jx) / (float)W) - 1.0f;
  float ny = 1.0f - 2.0f * ((py + jy) / (float)H);
  float3 d = f3(cam.f[0] + nx * cam.r[0] + ny * cam.u[0],
                cam.f[1] + nx * cam.r[1] + ny * cam.u[1],
                cam.f[2] + nx * cam.r[2] + ny * cam.u[2]);
  return normalize(d);
}

__global__ void k_forward(DevScene sc, DevCam cam, int H, int W, int spp,
                          int B, int mode, unsigned long long seed,
                          float* img) {
  int pid = blockIdx.x * blockDim.x + threadIdx.x;
  if (pid >= H * W) return;
  int px = pid % W, py = pid / W;
  float3 ro = f3(cam.p[0], cam.p[1], cam.p[2]);
  float3 acc = f3(0.0f, 0.0f, 0.0f);
  GCtx g = {};
  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);
    float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
    trace_path<false>(sc, ro, rd, st, B, mode, acc, g);
  }
  float inv_spp = 1.0f / (float)spp;
  img[pid * 3 + 0] = acc.x * inv_spp;
  img[pid * 3 + 1] = acc.y * inv_spp;
  img[pid * 3 + 2] = acc.z * inv_spp;
}

__global__ void k_backward(DevScene sc, DevCam cam, int H, int W, int spp,
                           int B, int mode, unsigned long long seed,
                           const float* gimg, float* g_tex, float* g_etex,
                           float* g_env, float* g_med) {
  extern __shared__ float smem[];
  bool staged_t = sc.nt <= kSharedTexels;
  bool staged_e = sc.net <= kSharedEmiTexels;
  float* s_gt = staged_t ? smem : nullptr;
  int off = staged_t ? sc.nt * 3 : 0;
  float* s_get = staged_e ? smem + off : nullptr;
  if (staged_e) off += sc.net * 3;
  float* s_gm = smem + off;  // [6]
  int total = off + 6;
  for (int i = threadIdx.x; i < total; i += blockDim.x) smem[i] = 0.0f;
  __syncthreads();

  int pid = blockIdx.x * blockDim.x + threadIdx.x;
  if (pid < H * W) {
    int px = pid % W, py = pid / W;
    float inv_spp = 1.0f / (float)spp;
    float gs[3] = {gimg[pid * 3 + 0] * inv_spp, gimg[pid * 3 + 1] * inv_spp,
                   gimg[pid * 3 + 2] * inv_spp};
    GCtx g;
    g.gs = gs;
    g.g_tex = g_tex;
    g.s_gt = s_gt;
    g.g_etex = g_etex;
    g.s_get = s_get;
    g.g_env = g_env;
    g.s_gm = sc.has_med ? s_gm : nullptr;
    float3 ro = f3(cam.p[0], cam.p[1], cam.p[2]);
    float3 acc = f3(0.0f, 0.0f, 0.0f);
    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);
      float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
      trace_path<true>(sc, ro, rd, st, B, mode, acc, g);
    }
  }

  __syncthreads();
  if (staged_t) {
    for (int i = threadIdx.x; i < sc.nt * 3; i += blockDim.x)
      if (s_gt[i] != 0.0f) atomicAdd(&g_tex[i], s_gt[i]);
  }
  if (staged_e) {
    for (int i = threadIdx.x; i < sc.net * 3; i += blockDim.x)
      if (s_get[i] != 0.0f) atomicAdd(&g_etex[i], s_get[i]);
  }
  if (sc.has_med) {
    for (int i = threadIdx.x; i < 6; i += blockDim.x)
      if (s_gm[i] != 0.0f) atomicAdd(&g_med[i], s_gm[i]);
  }
}

DevScene make_scene(const PtdSceneArgs& a) {
  DevScene sc;
  sc.tris = a.tris;
  sc.mats = a.mat_ids;
  sc.uvs = a.uvs;
  sc.nf = a.nodes_f;
  sc.ni = a.nodes_i;
  sc.n_nodes = a.n_nodes;
  sc.lf = a.light_faces;
  sc.lcdf = a.light_cdf;
  sc.nl = a.n_lights;
  sc.larea = a.total_light_area;
  sc.tex = a.tex;
  sc.thdr = a.tex_hdr;
  sc.nt = a.n_texels;
  sc.etex = a.emi_tex;
  sc.ehdr = a.emi_hdr;
  sc.net = a.n_emi_texels;
  sc.mtype = a.mat_type;
  sc.mrough = a.mat_rough;
  sc.mior = a.mat_ior;
  sc.nm = a.n_mats;
  sc.med_sa = a.med_sa;
  sc.med_ss = a.med_ss;
  sc.med_sbar = a.med_sbar;
  sc.has_med = a.has_med;
  sc.env = a.env;
  sc.ew = a.env_w;
  sc.eh = a.env_h;
  sc.ecdf_m = a.env_cdf_m;
  sc.ecdf_c = a.env_cdf_c;
  sc.epdf = a.env_pdf;
  return sc;
}

DevCam make_cam(const float* cam) {
  DevCam c;
  for (int i = 0; i < 3; ++i) {
    c.p[i] = cam[i];
    c.f[i] = cam[3 + i];
    c.r[i] = cam[6 + i];
    c.u[i] = cam[9 + i];
  }
  return c;
}

}  // namespace

extern "C" void ptd_forward_launch(const PtdSceneArgs* args, const float* cam,
                                   int H, int W, int spp, int max_bounces,
                                   int mode, long long seed, float* image,
                                   cudaStream_t stream) {
  DevScene sc = make_scene(*args);
  DevCam dc = make_cam(cam);
  int blocks = (H * W + kThreads - 1) / kThreads;
  k_forward<<<blocks, kThreads, 0, stream>>>(
      sc, dc, H, W, spp, max_bounces, mode, (unsigned long long)seed, image);
}

extern "C" void ptd_backward_launch(const PtdSceneArgs* args, const float* cam,
                                    int H, int W, int spp, int max_bounces,
                                    int mode, long long seed,
                                    const float* grad_image, float* grad_tex,
                                    float* grad_emi_tex, float* grad_env,
                                    float* grad_med, cudaStream_t stream) {
  DevScene sc = make_scene(*args);
  DevCam dc = make_cam(cam);
  int blocks = (H * W + kThreads - 1) / kThreads;
  size_t smem = 6 * sizeof(float);
  if (args->n_texels <= kSharedTexels)
    smem += (size_t)args->n_texels * 3 * sizeof(float);
  if (args->n_emi_texels <= kSharedEmiTexels)
    smem += (size_t)args->n_emi_texels * 3 * sizeof(float);
  k_backward<<<blocks, kThreads, smem, stream>>>(
      sc, dc, H, W, spp, max_bounces, mode, (unsigned long long)seed,
      grad_image, grad_tex, grad_emi_tex, grad_env, grad_med);
}