#version 460
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_EXT_shader_16bit_storage : require

// Ray Reconstruction hkpn output stage, preset E (no diamond-wallaby branch).
// One workgroup per 16x16 output tile. Stage 1 stages the render-space neighbourhood of
// the tile into shared memory (colour log, per-texel weights, anisotropy); stage 2 gathers
// a 5x5 elliptical kernel per output pixel and writes the two output surfaces.
// Parameter block (binding 0) holds NVIDIA's kernel argument block; word index = byte / 4.

layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;

layout(std430, binding = 0) readonly buffer Params { uint P[]; };
layout(std430, binding = 2) buffer Arena { uint A[]; };
layout(binding = 3) uniform sampler2D colourTex;   // arg 176, 640x360 rgba16f
layout(binding = 5) uniform sampler2D exposureTex; // arg 200, 1x1 rgba16f
layout(binding = 8) uniform sampler2D albedoA;     // arg 280, 640x360 rgba8unorm
layout(binding = 9) uniform sampler2D albedoB;     // arg 288, 640x360 rgba8unorm
layout(binding = 11) uniform sampler2D decTex;     // arg 448, 640x384 rgba16f
layout(rgba16f, binding = 12) uniform writeonly image2D outLog; // arg 224
layout(rgba16f, binding = 13) uniform writeonly image2D outLin; // arg 240

#define F_LOG2E    uintBitsToFloat(0x3FB8AA3Bu)
#define F_2LOG2E   uintBitsToFloat(0x4038AA3Bu)
#define F_LOG10_2  uintBitsToFloat(0x3E9A209Bu)
#define F_0_36     uintBitsToFloat(0x3EB851ECu)
#define F_100      uintBitsToFloat(0x42C80000u)
#define F_1E4      uintBitsToFloat(0x38D1B717u)
#define F_0_005    uintBitsToFloat(0x3BA3D70Au)
#define F_0_1      uintBitsToFloat(0x3DCCCCCDu)
#define F_1_7      uintBitsToFloat(0x3FD9999Au)
#define F_0_17     uintBitsToFloat(0x3E2E147Bu)
#define F_0_95     uintBitsToFloat(0x3F733333u)
#define F_0_6      uintBitsToFloat(0x3F19999Au)
#define F_9        uintBitsToFloat(0x41102CB4u)

// shared tile: 21 entries per row, up to 21x21 entries (render scale <= 1)
#define TILE_STRIDE 21
// binary16 numbers, stored as such: half the shared-memory traffic of the 5x5 gather
shared f16vec4 tLog[TILE_STRIDE * TILE_STRIDE]; // (L_R, L_G, L_B, alpha = 0) log-domain colour
shared f16vec4 tSig[TILE_STRIDE * TILE_STRIDE]; // (s0, s1, s2, s3) sigmoid of arena bytes 0..3
shared vec4 tAn[TILE_STRIDE * TILE_STRIDE];  // (p0, p1, p2, p3) anisotropy parameters

float pf(uint idx) { return uintBitsToFloat(P[idx]); }

// e4m3 byte -> float (exact); e4m3fn NaN for 0x7F/0xFF.
float e4m3(uint b) {
    uint e = (b >> 3u) & 15u;
    uint m = b & 7u;
    if (e == 15u && m == 7u) return uintBitsToFloat(0x7FC00000u);
    float v = (e == 0u) ? float(m) * uintBitsToFloat(0x3B000000u)
                        : (1.0 + float(m) * 0.125) * exp2(float(e) - 7.0);
    return ((b & 128u) != 0u) ? -v : v;
}

// round to nearest even binary16, returned as f32: the hardware conversion (RNE), kept from being folded back to x by
// OR-ing an opaque zero into its bits (RADV folds a register-only f32 -> f16 -> f32)
// Beyond the binary16 range NVIDIA's arithmetic keeps a finite rounded value: that case takes the explicit rounding.
uint16_t zero16;
float f16rnSlow(float x) {
    if (x == 0.0 || isnan(x) || isinf(x)) return x;
    int e;
    frexp(x, e);
    int ee = max(e - 11, -24);
    float q = uintBitsToFloat(uint(ee + 127) << 23);
    return roundEven(x / q) * q;
}
float f16rn(float x) {
    if (abs(x) >= 65504.0) return f16rnSlow(x);
    return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | zero16));
}

// Sigmoid as evaluated in binary16 on this lowering: r = f16(2^(-x*log2e)), den = f16(1+r),
// s = f16(1/den), with one Newton step for small s and the observed fix-ups of r.
float sig16(float e) {
    float x = f16rn(e);
    float r = f16rn(exp2(-x * F_LOG2E));
    uint hb = uint(float16BitsToUint16(float16_t(r)));
    if (hb == 0x1F79u || hb == 0x25CFu) {
        r = f16rn(r - 0.0009765625);
    } else if (hb == 0xC13Bu) {
        r = f16rn(r + 0.00006103515625);
    } else if (hb == 0xC1EFu) {
        r = f16rn(r + 0.000030517578125);
    }
    float d32 = f16rn(1.0 + r);
    float rc = 1.0 / d32;
    float s16 = f16rn(rc);
    uint sb = uint(float16BitsToUint16(float16_t(s16))) & 0x7FFFu;
    if (sb != 0u && sb < 0x8Fu) {
        float ee = fma(-d32, s16, 1.0);
        s16 = f16rn(fma(rc, ee, s16));
    }
    return s16;
}

// tanh as evaluated: polynomial for small |x|, exponential form otherwise.
float tanhApprox(float x) {
    float ax = abs(x);
    float x2 = x * x;
    float q = fma(uintBitsToFloat(0x3C80F082u), x2, uintBitsToFloat(0xBD563CAEu));
    q = fma(q, x2, uintBitsToFloat(0x3E085941u));
    q = fma(q, x2, uintBitsToFloat(0xBEAAA9EDu));
    q = fma(q, x2, 0.0);
    float small = fma(q, x, x);
    float ex = exp2(ax * F_2LOG2E);
    float big = fma(1.0 / (ex + 1.0), -2.0, 1.0);
    if (ax >= F_9) big = 1.0;
    float bigSigned = uintBitsToFloat((floatBitsToUint(x) & 0x80000000u) | floatBitsToUint(big));
    return (ax >= F_0_6) ? bigSigned : small;
}

// round half away from zero
int roundAway(float v) {
    return int(trunc(v + ((v < 0.0) ? -0.5 : 0.5)));
}

void main() {
    zero16 = uint16_t(P[0] >> 24u);   // the output width is below 2^24
    // ---- parameters (byte offsets / 4) ----
    uint outW = P[0];
    uint outH = P[1];
    float sx = pf(6u);     // 24: render/out scale x
    float sy = pf(7u);     // 28
    float ox = pf(10u);    // 40: jitter offset x
    float oy = pf(11u);    // 44
    float expScale = pf(37u); // 148
    uint decMaxX = P[14] - 1u; // 56: dec0 width
    uint decMaxY = P[15] - 1u; // 60: dec0 height
    uint colMaxX = P[35];      // 140: colour/albedo max x
    uint colMaxY = P[36];      // 144: colour/albedo max y
    uint arenaBase = P[64];    // 256: arena offset
    uint arenaPitch = P[4];    // 16: render width (640)

    // ---- tile origin / size in render space ----
    uvec2 wg = gl_WorkGroupID.xy;
    float tileX = ox + fma(float(wg.x * 16u) + 0.5, sx, -0.5);
    float tileY = oy + fma(float(wg.y * 16u) + 0.5, sy, -0.5);
    int c0 = int(floor(tileX)) - 2;
    int r0 = int(floor(tileY)) - 2;
    int tw = int(ceil(16.0 * sx)) + 5;
    int th = int(ceil(16.0 * sy)) + 5;

    // ---- stage 1: tile load ----
    uint total = uint(tw * th);
    for (uint idx = gl_LocalInvocationIndex; idx < total; idx += 256u) {
        int col = int(idx) % tw;
        int row = int(idx) / tw;
        int ix = clamp(c0 + col, 0, int(decMaxX));
        int iy = clamp(r0 + row, 0, int(decMaxY));

        vec3 d = texelFetch(decTex, ivec2(ix, iy), 0).rgb;
        int ax = min(ix, int(colMaxX));
        int ay = min(iy, int(colMaxY));
        vec3 alb = texelFetch(albedoA, ivec2(ax, ay), 0).rgb
                 + texelFetch(albedoB, ivec2(ax, ay), 0).rgb;
        float albSum = (alb.r + alb.g) + alb.b;
        if (albSum != 0.0) {
            d.r *= max(alb.r, F_0_005);
            d.g *= max(alb.g, F_0_005);
            d.b *= max(alb.b, F_0_005);
        }
        vec3 lg;
        lg.r = f16rn(F_0_36 * (log2(max(fma(d.r, F_100, 1.0), F_1E4)) * F_LOG10_2));
        lg.g = f16rn(F_0_36 * (log2(max(fma(d.g, F_100, 1.0), F_1E4)) * F_LOG10_2));
        lg.b = f16rn(F_0_36 * (log2(max(fma(d.b, F_100, 1.0), F_1E4)) * F_LOG10_2));

        uint wordIdx = (arenaBase >> 2u) + uint(arenaPitch * uint(iy) + uint(ix)) * 2u;
        uint lo = A[wordIdx];
        uint hi = A[wordIdx + 1u];
        vec4 sg = vec4(sig16(e4m3(lo & 0xFFu)),
                       sig16(e4m3((lo >> 8u) & 0xFFu)),
                       sig16(e4m3((lo >> 16u) & 0xFFu)),
                       sig16(e4m3(lo >> 24u)));
        vec4 an = vec4(e4m3(hi & 0xFFu),
                       e4m3((hi >> 8u) & 0xFFu),
                       e4m3((hi >> 16u) & 0xFFu),
                       e4m3(hi >> 24u));

        uint se = uint(row * TILE_STRIDE + col);
        tLog[se] = f16vec4(lg, 0.0);
        tSig[se] = f16vec4(sg);
        tAn[se] = an;
    }
    barrier();

    // ---- stage 2: per output pixel ----
    uvec2 pix = gl_GlobalInvocationID.xy;
    if (pix.x >= outW || pix.y >= outH) return;

    float expo = texelFetch(exposureTex, ivec2(0, 0), 0).r;
    float e41 = ((expo == 0.0) ? 1.0 : expo) * expScale;

    float fx = (float(pix.x) + 0.5) * sx;
    float fy = (float(pix.y) + 0.5) * sy;
    float cx = fx + ox;
    float cy = fy + oy;

    // colour fetch: log of exposure-scaled colour at the nearest render texel
    int ux = clamp(roundAway((fx - 0.5) + ox), 0, int(colMaxX));
    int uy = clamp(roundAway((fy - 0.5) + oy), 0, int(colMaxY));
    vec3 col = texelFetch(colourTex, ivec2(ux, uy), 0).rgb;
    vec3 gcol;
    gcol.r = F_0_36 * (log2(max(fma(col.r * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
    gcol.g = F_0_36 * (log2(max(fma(col.g * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
    gcol.b = F_0_36 * (log2(max(fma(col.b * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);

    // centre entry, anisotropy
    float flx = floor(cx);
    float fly = floor(cy);
    int ix0 = int(flx) - c0;
    int iy0 = int(fly) - r0;
    float dx0 = (flx + 0.5) - cx;
    float dy0 = (fly + 0.5) - cy;
    float ks = (sy >= 0.5) ? 1.0 : (0.5 / sy);

    vec4 an = tAn[iy0 * TILE_STRIDE + ix0];
    float sa = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.x)), F_1_7, F_0_17);
    float sb = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.y)), F_1_7, F_0_17);
    float rho = F_0_95 * tanhApprox(an.z);

    float sa2 = sa * sa;
    float sb2 = sb * sb;
    float denom = (sa2 * sb2) * (1.0 - rho * rho);
    float invDenom = -(1.0 / denom);
    float cross2 = rho * (2.0 * sqrt(sa2 * sb2));

    // 5x5 gather
    float sumw = 0.0;
    vec3 accL = vec3(0.0);
    vec4 accS = vec4(0.0);
    for (int c = -2; c <= 2; c++) {
        for (int r = -2; r <= 2; r++) {
            float dx = ks * (dx0 + float(c));
            float dy = ks * (dy0 + float(r));
            uint se = uint((iy0 + r) * TILE_STRIDE + (ix0 + c));
            precise float t_c = (cross2 * dx) * dy;
            precise float t_b = dx * (sb2 * dx);
            precise float t_a = dy * (sa2 * dy);
            precise float num = (t_c - t_b) - t_a;
            precise float Q = invDenom * num;
            float w = exp2((Q * -0.5) * F_LOG2E);
            sumw += w;
            vec4 lv = vec4(tLog[se]);
            vec4 sv = vec4(tSig[se]);
            accL = fma(vec3(w), lv.rgb, accL);
            accS = fma(vec4(w), sv, accS);
        }
    }

    float inv = 1.0 / sumw;
    vec3 Lm = accL * inv;
    float S0 = accS.x * inv;
    float S1 = accS.y * inv;
    float S2 = accS.z * inv;

    // blend weight
    float q = F_0_1 / (sx * sy);
    float t1 = sumw * q;
    float t2 = S0 * t1;
    float t3 = (sx * sy) * t2;
    float t4 = 4.0 * t3;
    float t5 = (sy < 0.5) ? t4 : t2;
    float wgt = (t5 < 0.0) ? 0.0 : ((t5 > 1.0) ? 1.0 : t5);

    float fA = fma(1.0 - S1, wgt, S1);
    float fB = fma(1.0 - S2, wgt, S2);

    // surface 224 (E): fA * D + (1 - fA) * gcol; D = Lm (f16 tile log channels),
    // gcol = log of the exposure-scaled colour texel (PTX linear texel fetch path); D.a = 0 so aA = 1 - fA
    vec3 cA = fma(Lm, vec3(fA), gcol * (1.0 - fA));
    float aA = 1.0 - fA;
    float sumA = cA.r + cA.g + cA.b + aA;
    if (isnan(sumA) || isinf(sumA)) {
        cA = Lm;
        aA = 0.0;
    }
    imageStore(outLog, ivec2(pix), vec4(f16rn(cA.r), f16rn(cA.g), f16rn(cA.b), f16rn(aA)));

    // surface 240 (E, PTX BB0_83): {R,G,B,1} of lin = (10^(clamp(c)/0.36)-1)/100/e41,
    // c = fma(Lm, fB, gcol*(1-fB)); c = Lm when the sum is not finite; stored RNE f16
    vec3 cB = fma(Lm, vec3(fB), gcol * (1.0 - fB));
    float sumB = cB.r + cB.g + cB.b;
    if (isnan(sumB) || isinf(sumB)) cB = Lm;
    cB = clamp(cB, vec3(F_1E4), vec3(3.0));
    // PTX divides by constants and takes lg2(10) with the hardware approximations at run time; keep the
    // constants opaque so the compiler cannot fold them with different rounding.
    uint zero = P[0] >> 24u;
    float c036 = uintBitsToFloat(0x3EB851ECu | zero);
    float c100 = uintBitsToFloat(0x42C80000u | zero);
    float lg10 = log2(uintBitsToFloat(0x41200000u | zero));
    float rc036 = 1.0 / c036;
    float rc100 = 1.0 / c100;
    float rce41 = 1.0 / e41;
    precise vec3 lin = (exp2((cB * rc036) * lg10) - 1.0) * rc100 * rce41;
    imageStore(outLin, ivec2(pix), vec4(float(float16_t(lin.r)), float(float16_t(lin.g)), float(float16_t(lin.b)), 1.0));
}