// daedalus-fourier — H.264 4x4 inverse integer transform + add, V3D 7.1.
//
// H.264 spec §8.5.12.1.  Pure integer arithmetic — no trig constants
// (unlike VP9 IDCT 8x8).  Row pass first, column pass second; round
// (+32) >> 6, add to dst, clip to u8.
//
// Block memory layout: COLUMN-MAJOR.  block[c*4 + r] = coefficient at
// (row r, column c).  Matches FFmpeg `ff_h264_idct_add_neon`.
//
// Workgroup layout: 64 invocations = 4 lanes/block × 16 blocks/WG.
//   - row pass: lane k (0..3) reads row k of the block (4 coefficients,
//               one from each column), runs the butterfly, writes 4
//               outputs to one row of tmp_shared.
//   - column pass: lane k reads column k of tmp_shared (4 rows),
//                  runs the butterfly, writes 4 outputs to dst as
//                  column k at rows 0..3.
//
// shared = 16 × 16 × 4 B = 1 KiB.  Well under V3D's 16 KiB limit.
//
// License: BSD-2-Clause.

#version 450
#extension GL_EXT_shader_8bit_storage             : require
#extension GL_EXT_shader_16bit_storage            : require
#extension GL_EXT_shader_explicit_arithmetic_types : require

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

layout(binding = 0) readonly buffer Coeffs {
    int16_t coeffs[];   // N × 16 column-major
} u_coeffs;

layout(binding = 1) buffer Dst {
    uint8_t dst[];      // H × stride bytes (caller-provided base)
} u_dst;

layout(binding = 2) readonly buffer Meta {
    uvec4 meta[];       // .x = dst_off (byte offset into u_dst.dst)
} u_meta;

layout(push_constant) uniform PC {
    uint n_blocks;
    uint dst_stride_u8;
    uint _pad0, _pad1;
} pc;

// 16 blocks per WG × 16 ints per block = 256 ints = 1 KiB shared.
shared int tmp_shared[16 * 16];

// 1D butterfly per H.264 §8.5.12.1.  d[0..3] in, o[0..3] out.
void idct4_1d(int d0, int d1, int d2, int d3,
              out int o0, out int o1, out int o2, out int o3)
{
    int e = d0 + d2;
    int f = d0 - d2;
    int g = (d1 >> 1) - d3;
    int h = d1 + (d3 >> 1);
    o0 = e + h;
    o1 = f + g;
    o2 = f - g;
    o3 = e - h;
}

void main()
{
    // Lane decomposition: local_size 64 = 16 blocks × 4 lanes/block.
    uint gid          = gl_GlobalInvocationID.x;
    uint wg_id        = gid / 64u;
    uint lane_in_wg   = gid & 63u;
    uint block_local  = lane_in_wg >> 2;          // 0..15
    uint k            = lane_in_wg & 3u;          // 0..3
    uint block_idx    = wg_id * 16u + block_local;

    bool oob = (block_idx >= pc.n_blocks);

    // ---- Row pass --------------------------------------------------
    // lane k handles row r=k.  Reads block[c*4 + k] for c=0..3 (one
    // element from each column at fixed row).
    if (!oob) {
        uint base = block_idx * 16u;
        int d0 = int(u_coeffs.coeffs[base + 0u * 4u + k]);
        int d1 = int(u_coeffs.coeffs[base + 1u * 4u + k]);
        int d2 = int(u_coeffs.coeffs[base + 2u * 4u + k]);
        int d3 = int(u_coeffs.coeffs[base + 3u * 4u + k]);

        int o0, o1, o2, o3;
        idct4_1d(d0, d1, d2, d3, o0, o1, o2, o3);

        // Write row k of tmp_shared[block_local].
        uint tbase = block_local * 16u + k * 4u;
        tmp_shared[tbase + 0u] = o0;
        tmp_shared[tbase + 1u] = o1;
        tmp_shared[tbase + 2u] = o2;
        tmp_shared[tbase + 3u] = o3;
    }

    barrier();

    // ---- Column pass ----------------------------------------------
    // lane k handles column c=k.  Reads tmp[r][k] for r=0..3.
    if (!oob) {
        uint tbase = block_local * 16u;
        int s0 = tmp_shared[tbase + 0u * 4u + k];
        int s1 = tmp_shared[tbase + 1u * 4u + k];
        int s2 = tmp_shared[tbase + 2u * 4u + k];
        int s3 = tmp_shared[tbase + 3u * 4u + k];

        int o0, o1, o2, o3;
        idct4_1d(s0, s1, s2, s3, o0, o1, o2, o3);

        // Column k at rows 0..3 of dst, offset by meta.x (dst_off).
        uint dst_off = u_meta.meta[block_idx].x;
        uint stride  = pc.dst_stride_u8;
        uint a0 = dst_off + 0u * stride + k;
        uint a1 = dst_off + 1u * stride + k;
        uint a2 = dst_off + 2u * stride + k;
        uint a3 = dst_off + 3u * stride + k;

        int p0 = int(u_dst.dst[a0]);
        int p1 = int(u_dst.dst[a1]);
        int p2 = int(u_dst.dst[a2]);
        int p3 = int(u_dst.dst[a3]);

        u_dst.dst[a0] = uint8_t(clamp(p0 + ((o0 + 32) >> 6), 0, 255));
        u_dst.dst[a1] = uint8_t(clamp(p1 + ((o1 + 32) >> 6), 0, 255));
        u_dst.dst[a2] = uint8_t(clamp(p2 + ((o2 + 32) >> 6), 0, 255));
        u_dst.dst[a3] = uint8_t(clamp(p3 + ((o3 + 32) >> 6), 0, 255));
    }
}