#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <math_constants.h>
#include <stdio.h>
#include <math.h>
#include "tile_backward_fp16.h"
#include "tile_common_fp16.h"

#ifdef __CUDACC__
    #pragma STDC FP_CONTRACT OFF
#endif

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) (x)
#endif

#define DEBUG_CUDA_KERNELS_FP16_BACKWARD 0
#define DEBUG_SEQUENTIAL_FP16 0  // Set to 1 for sequential debugging, 0 for parallel execution
#define SIZE_BASED_OPTIM 0  // Set to 1 to use size-based c_blend and rotation (image_width * SIZE_THRESHOLD_RATIO threshold), 0 for original global c_blend
#define SIZE_THRESHOLD_RATIO 0.2f  // Threshold ratio for size-based optimization (image_width * SIZE_THRESHOLD_RATIO)
#define C_BLEND_INVERSE 0  // Set to 1 to apply c_blend to large primitives (s > threshold), 0 to apply to small primitives (s <= threshold)

// Color map sampling function (same as tile_forward_fp16.cu)
__device__ __forceinline__ float3 sample_color_map_fp16(
    const __half* color_map, int map_h, int map_w, 
    float x, float y, float radius, float rotation) {
    
    // Transform coordinates to color map space
    // 1. Normalize by primitive radius to get [-1, 1] range
    float r_inv = radius > 1e-6f ? (1.0f / radius) : 1e6f;
    float norm_x = x * r_inv;
    float norm_y = y * r_inv;
    
    // 2. Rotate
    float cos_theta = cosf(rotation);
    float sin_theta = sinf(rotation);
    float rot_x =  cos_theta * norm_x + sin_theta * norm_y;
    float rot_y = -sin_theta * norm_x + cos_theta * norm_y;
    
    // 3. Convert normalized [-1, 1] coordinates to color map pixel coordinates [0, W-1], [0, H-1]
    float coord_x = (rot_x + 1.0f) * 0.5f * (map_w - 1);
    float coord_y = (rot_y + 1.0f) * 0.5f * (map_h - 1);
    
    // Clamp coordinates
    coord_x = fmaxf(0.0f, fminf(map_w - 1, coord_x));
    coord_y = fmaxf(0.0f, fminf(map_h - 1, coord_y));
    
    // Bilinear interpolation
    int x0 = (int)coord_x;
    int y0 = (int)coord_y;
    int x1 = min(x0 + 1, map_w - 1);
    int y1 = min(y0 + 1, map_h - 1);
    
    float fx = coord_x - x0;
    float fy = coord_y - y0;
    
    // Sample colors (convert from __half to float)
    float3 c00 = make_float3(
        __half2float(color_map[3 * (y0 * map_w + x0) + 0]),
        __half2float(color_map[3 * (y0 * map_w + x0) + 1]),
        __half2float(color_map[3 * (y0 * map_w + x0) + 2])
    );
    float3 c01 = make_float3(
        __half2float(color_map[3 * (y1 * map_w + x0) + 0]),
        __half2float(color_map[3 * (y1 * map_w + x0) + 1]),
        __half2float(color_map[3 * (y1 * map_w + x0) + 2])
    );
    float3 c10 = make_float3(
        __half2float(color_map[3 * (y0 * map_w + x1) + 0]),
        __half2float(color_map[3 * (y0 * map_w + x1) + 1]),
        __half2float(color_map[3 * (y0 * map_w + x1) + 2])
    );
    float3 c11 = make_float3(
        __half2float(color_map[3 * (y1 * map_w + x1) + 0]),
        __half2float(color_map[3 * (y1 * map_w + x1) + 1]),
        __half2float(color_map[3 * (y1 * map_w + x1) + 2])
    );
    
    // Interpolate
    float3 c0 = make_float3(
        c00.x * (1.0f - fx) + c10.x * fx,
        c00.y * (1.0f - fx) + c10.y * fx,
        c00.z * (1.0f - fx) + c10.z * fx
    );
    float3 c1 = make_float3(
        c01.x * (1.0f - fx) + c11.x * fx,
        c01.y * (1.0f - fx) + c11.y * fx,
        c01.z * (1.0f - fx) + c11.z * fx
    );
    
    return make_float3(
        c0.x * (1.0f - fy) + c1.x * fy,
        c0.y * (1.0f - fy) + c1.y * fy,
        c0.z * (1.0f - fy) + c1.z * fy
    );
}

__device__ __half2* d_grad_mxy_h2 = nullptr; // (mx,my)
__device__ __half2* d_grad_or_h2 = nullptr; // (opacity, radius)
__device__ __half2* d_grad_theta_cr_h2 = nullptr; // (rotation, color_r)
__device__ __half2* d_grad_cgb_h2 = nullptr; // (color_g, color_b)

// Host setter to inject device pointers without changing kernel ABI.
void SetPackedGradPtrs(__half2* mxy_h2, __half2* or_h2, __half2* theta_cr_h2, __half2* cgb_h2)
{
    cudaMemcpyToSymbol(d_grad_mxy_h2,      &mxy_h2,      sizeof(__half2*));
    cudaMemcpyToSymbol(d_grad_or_h2,       &or_h2,       sizeof(__half2*));
    cudaMemcpyToSymbol(d_grad_theta_cr_h2, &theta_cr_h2, sizeof(__half2*));
    cudaMemcpyToSymbol(d_grad_cgb_h2,      &cgb_h2,      sizeof(__half2*));
}

// ============================
// Unpack packed half2 buffers into legacy OutputTensorsFP16
// ============================
__global__ void unpack_packed_grads_kernel(
    int N,
    const __half2* __restrict__ grad_mxy_h2,       // (mx,my)
    const __half2* __restrict__ grad_or_h2,        // (opacity, radius)
    const __half2* __restrict__ grad_theta_cr_h2,  // (rotation, color_r)
    const __half2* __restrict__ grad_cgb_h2,       // (color_g, color_b)
    __half* __restrict__ grad_means2D,             // [2N]
    __half* __restrict__ grad_opacities,           // [N]
    __half* __restrict__ grad_radii,               // [N]
    __half* __restrict__ grad_rotations,           // [N]
    __half* __restrict__ grad_colors               // [3N] interleaved RGB
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;

    // means2D
    __half2 mxy = grad_mxy_h2[n];
    grad_means2D[2*n + 0] = __low2half(mxy);
    grad_means2D[2*n + 1] = __high2half(mxy);

    // opacity, radius
    __half2 orv = grad_or_h2[n];
    grad_opacities[n] = __low2half(orv);
    grad_radii[n]     = __high2half(orv);

    // rotation, color_r
    __half2 th_cr = grad_theta_cr_h2[n];
    grad_rotations[n]    = __low2half(th_cr);
    grad_colors[3*n + 0] = __high2half(th_cr);

    // color_g, color_b
    __half2 gb = grad_cgb_h2[n];
    grad_colors[3*n + 1] = __low2half(gb);
    grad_colors[3*n + 2] = __high2half(gb);
}

void UnpackPackedGrads(
    int N,
    const __half2* grad_mxy_h2, const __half2* grad_or_h2, const __half2* grad_theta_cr_h2, const __half2* grad_cgb_h2, 
    const OutputTensorsFP16& out,
    cudaStream_t stream = 0)
{
    int bs = 256, gs = (N + bs - 1) / bs;
    unpack_packed_grads_kernel<<<gs, bs, 0, stream>>>(
        N,
        grad_mxy_h2,
        grad_or_h2,
        grad_theta_cr_h2,
        grad_cgb_h2,
        out.grad_means2D,
        out.grad_opacities,
        out.grad_radii,
        out.grad_rotations,
        out.grad_colors
    );
}

// Per-pixel backward with OVER rule and per-pixel local-k mapping
template<bool HAVE_ALPHA_CACHE, bool HAVE_T_CACHE>
__device__ inline void backward_over_one_pixel_fp16(
    // pixel coord & config
    int x, int y, const TileConfigFP16 tile_config,
    const PrimitiveConfigFP16 prim_config,
    // primitive set for this tile
    const int* __restrict__ prim_ids,  // [K]
    int K,
    // upstream grads (dL/dC and dL/dA at this pixel)
    __half gCx, __half gCy, __half gCz, __half gA,
    // structured IO
    const InputTensorsFP16 inputs,
    const OutputTensorsFP16 outputs,
    const GlobalBuffersFP16 buffers,
    const LearningRateConfigFP16 lr_config
){
    const int W = tile_config.image_width;
    const int H = tile_config.image_height;
    const int pixel_idx = y * W + x;
    const int base = pixel_idx * prim_config.max_prims_per_pixel;
    const int num_k = min(buffers.pixel_prim_counts[pixel_idx], prim_config.max_prims_per_pixel);
    if (num_k <= 0) return;

    // Rebuild local k by using the SAME push condition as forward (radius check only).
    int k_running = 0;

    for (int kk = 0; kk < K; ++kk){
        const int n = prim_ids[kk];
        if (n < 0 || n >= prim_config.num_primitives) continue;

        // Forward push condition: radius check only
        const __half mx = inputs.means2D[2*n+0];
        const __half my = inputs.means2D[2*n+1];
        const __half r  = inputs.radii[n];
        const __half dx = __hsub(__float2half((float)x), mx);
        const __half dy = __hsub(__float2half((float)y), my);
        const int template_idx = inputs.global_bmp_sel[n];
        if (template_idx >= 0 && template_idx < prim_config.num_templates) {
            const __half template_width = __float2half((float)prim_config.template_width);
            const __half template_height = __float2half((float)prim_config.template_height);
            const __half max_dim = __hgt(template_width, template_height) ? template_width : template_height;
            const __half scale_factor = __hdiv(r, __hmul(max_dim, H05));
            
            const __half half_width = __hmul(__hmul(template_width, H05), scale_factor);
            const __half half_height = __hmul(__hmul(template_height, H05), scale_factor);
            
            if (__hgt(__habs(dx), half_width) || __hgt(__habs(dy), half_height)) {
                continue;
            }
        } else {
            if (__hgt(__hfma(dy, dy, __hmul(dx, dx)), __hmul(r, r))) {
                continue;
            }
        }

        const int k = k_running;
        if (k >= num_k) { k_running++; continue; } // safety

        // Read caches written by forward
        const int idx = base + k;
        const __half a_k = buffers.pixel_alphas[idx];
        const __half T_k = buffers.pixel_T_values[idx];
        const __half cr  = buffers.pixel_colors_r[idx];
        const __half cg  = buffers.pixel_colors_g[idx];
        const __half cb  = buffers.pixel_colors_b[idx];

        // Get mask value early for gradient computation decisions
        __half mask_val = H0;
#if SIZE_BASED_OPTIM
        // Calculate rotation_scale early for use in mask computation
        const float threshold = (float)tile_config.image_width * SIZE_THRESHOLD_RATIO;
        const float s_radius = __half2float(inputs.radii[n]);
        // const float c_blend_local = fmaxf(0.0f, fminf(1.0f, 1.0f - (s_radius / threshold)));
        // const float rotation_scale = 1.0f - c_blend_local;
#if C_BLEND_INVERSE
        // Binary c_blend: s > threshold -> inputs.c_blend, s <= threshold -> 0 (apply to large primitives)
        const float c_blend_local = (s_radius > threshold) ? inputs.c_blend : 0.0f;
        const float rotation_scale = (s_radius > threshold) ? 1.0f : 0.0f;
#else
        // Binary c_blend: s > threshold -> 0, s <= threshold -> inputs.c_blend (apply to small primitives)
        const float c_blend_local = (s_radius <= threshold) ? inputs.c_blend : 0.0f;
        const float rotation_scale = (s_radius <= threshold) ? 1.0f : 0.0f;
#endif
#endif
        if (template_idx >= 0 && template_idx < prim_config.num_templates) {
            const __half inv_r = __hgt(r, EPS1_FP16) ? __hdiv(H1, r) : EPS6_FP16;
#if SIZE_BASED_OPTIM
            const __half phi_original = inputs.rotations[n];
            const __half rotation_scale_h = __float2half(rotation_scale);
            const __half phi = __hmul(phi_original, rotation_scale_h);  // Apply size-based rotation scaling
#else
            const __half phi = inputs.rotations[n];
#endif
            const __half c = hcos(phi), s = hsin(phi);
            const __half ndx = __hmul(dx, inv_r);
            const __half ndy = __hmul(dy, inv_r);
            const __half u =  __hfma(s, ndy, __hmul(c, ndx));
            const __half v = __hfma(c, ndy, __hmul(__hneg(s), ndx));
            const __half tex_x = __hmul(__hfma(u, H05, H05), __float2half(prim_config.template_width  - 1));
            const __half tex_y = __hmul(__hfma(v, H05, H05), __float2half(prim_config.template_height - 1));

            const __half* tex = &inputs.primitive_templates[
                template_idx * prim_config.template_height * prim_config.template_width];
            mask_val = bilinear_sample_fp16(tex, prim_config.template_height, prim_config.template_width, tex_y, tex_x);
            mask_val = clampf_fp16(mask_val, H0, H1);
        }
        
        if (__hle(mask_val, EPS1_FP16)) {
            // This primitive doesn't affect this pixel, skip
            continue;
        }
            
        // Build suffix S and back-product B:
        // S = Σ_{m>k} c_m α_m T_m,  B = Π_{m>k}(1-α_m)
        __half Sx=H0, Sy=H0, Sz=H0, B=H1;
        for (int m = k+1; m < num_k; ++m){
            const int midx = base + m;
            const __half am = buffers.pixel_alphas[midx];
            const __half Tm = buffers.pixel_T_values[midx];
            const __half crm= buffers.pixel_colors_r[midx];
            const __half cgm= buffers.pixel_colors_g[midx];
            const __half cbm= buffers.pixel_colors_b[midx];
            // Fused multiply-add in half to reduce rounding
            Sx = __hfma(__hmul(crm, am), Tm, Sx);
            Sy = __hfma(__hmul(cgm, am), Tm, Sy);
            Sz = __hfma(__hmul(cbm, am), Tm, Sz);
            B  = __hmul(B, __hmax(__hsub(H1, am), H0));
        }

        // (1) dL/d(color logits): ∂L/∂c = gC * (α_k T_k); ∂c/∂z = σ(z)(1-σ(z))
        const __half sr = sigmoidf_safe_fp16(inputs.colors[3*n+0]);
        const __half sg = sigmoidf_safe_fp16(inputs.colors[3*n+1]);
        const __half sb = sigmoidf_safe_fp16(inputs.colors[3*n+2]);
        // Use FMA for gC * (a_k * T_k)
        const __half aT = __hmul(a_k, T_k);

        __half add_mx = H0;
        __half add_my = H0;
        __half add_opacity = H0;
        __half add_radius  = H0;
        __half add_rot = H0;      
        __half add_r = H0;
        __half add_g = H0;
        __half add_b = H0;
        
        // (2) dL/dα_k (OVER):
        // ∂C/∂α_k = T_k c_k - S/(1-α_k),  ∂A/∂α_k = T_k * B
        const __half inv1m = __hdiv(H1, __hmax(__hsub(H1, a_k), EPS1_FP16));
        const __half dCda_x = __hsub(__hmul(T_k, cr), __hmul(Sx, inv1m));
        const __half dCda_y = __hsub(__hmul(T_k, cg), __hmul(Sy, inv1m));
        const __half dCda_z = __hsub(__hmul(T_k, cb), __hmul(Sz, inv1m));
        // Accumulate dLdalpha with fused adds
        __half dLdalpha = __hfma(gCx, dCda_x, H0);
        dLdalpha = __hfma(gCy, dCda_y, dLdalpha);
        dLdalpha = __hfma(gCz, dCda_z, dLdalpha);
        dLdalpha = __hfma(gA, __hmul(T_k, B), dLdalpha);

        // (3) α = α_max * sigmoid(v_op) * mask(u,v)
        //     opacities are LOGITS  → dα/dv = α_max * mask * σ(v)*(1-σ(v))
        const __half v_op = inputs.opacities[n];
        const __half s_op = sigmoidf_safe_fp16(v_op);
        __half dmask_dx_tex=__float2half(0.0f), dmask_dy_tex=__float2half(0.0f);

        // Template coords (forward-identical): (u,v) from rotation φ and scale r
        const __half inv_r = __hgt(r, EPS1_FP16) ? __hdiv(H1, r) : EPS6_FP16;
#if SIZE_BASED_OPTIM
        const __half phi_original = inputs.rotations[n];
        const __half rotation_scale_h = __float2half(rotation_scale);
        const __half phi = __hmul(phi_original, rotation_scale_h);  // Apply size-based rotation scaling
#else
        const __half phi = inputs.rotations[n];
#endif
        const __half c = hcos(phi), s = hsin(phi);
        const __half ndx = __hmul(dx, inv_r);
        const __half ndy = __hmul(dy, inv_r);
        const __half u =  __hfma(s, ndy, __hmul(c, ndx));
        const __half v = __hfma(c, ndy, __hmul(__hneg(s), ndx));
        const __half tex_x = __hmul(__hfma(u, H05, H05), __float2half(prim_config.template_width  - 1));
        const __half tex_y = __hmul(__hfma(v, H05, H05), __float2half(prim_config.template_height - 1));

        if (template_idx >= 0 && template_idx < prim_config.num_templates) {
            const __half* tex = &inputs.primitive_templates[
                template_idx * prim_config.template_height * prim_config.template_width];
            mask_val = bilinear_value_and_grad_xy_fp16(
                tex, prim_config.template_height, prim_config.template_width,
                tex_y, tex_x, dmask_dx_tex, dmask_dy_tex);
            mask_val = clampf_fp16(mask_val, H0, H1);
        } else {
            mask_val = H0;
            dmask_dx_tex = H0;
            dmask_dy_tex = H0;
        }
        

        // Only compute gradients if mask value is significant (not transparent)
        if (__hgt(mask_val, EPS1_FP16)) {
            // colors
#if SIZE_BASED_OPTIM
            // Size-based c_blend and rotation: threshold = image_width * SIZE_THRESHOLD_RATIO
            const float threshold = (float)tile_config.image_width * SIZE_THRESHOLD_RATIO;
            const float s_radius = __half2float(inputs.radii[n]);
            // const float c_blend_local = fmaxf(0.0f, fminf(1.0f, 1.0f - (s_radius / threshold)));
            // const float rotation_scale = 1.0f - c_blend_local;  // Same as c_blend: large primitives get 0 rotation, small get full rotation
#if C_BLEND_INVERSE
            // Binary c_blend: s > threshold -> inputs.c_blend, s <= threshold -> 0 (apply to large primitives)
            const float c_blend_local = (s_radius > threshold) ? inputs.c_blend : 0.0f;
            const float rotation_scale = (s_radius > threshold) ? 1.0f : 0.0f;  // Large primitives get full rotation, small get 0 rotation
#else
            // Binary c_blend: s > threshold -> 0, s <= threshold -> inputs.c_blend (apply to small primitives)
            const float c_blend_local = (s_radius <= threshold) ? inputs.c_blend : 0.0f;
            const float rotation_scale = (s_radius <= threshold) ? 1.0f : 0.0f;  // Small primitives get full rotation, large get 0 rotation
#endif
            // c_i gradient: ∂L/∂c_i = gCx * (1 - c_blend) * (α_k T_k)
            const __half c_blend_h = __float2half(c_blend_local);
            const __half inv_blend = __hsub(H1, c_blend_h);
            add_r = __hmul(__hmul(__hmul(gCx, aT), inv_blend), __hmul(sr, __hsub(__float2half(1.0f), sr)));
            add_g = __hmul(__hmul(__hmul(gCy, aT), inv_blend), __hmul(sg, __hsub(__float2half(1.0f), sg)));
            add_b = __hmul(__hmul(__hmul(gCz, aT), inv_blend), __hmul(sb, __hsub(__float2half(1.0f), sb)));
#else
            // Original: no c_blend consideration for c_i gradient
            add_r = __hmul(__hmul(gCx, aT), __hmul(sr, __hsub(__float2half(1.0f), sr)));
            add_g = __hmul(__hmul(gCy, aT), __hmul(sg, __hsub(__float2half(1.0f), sg)));
            add_b = __hmul(__hmul(gCz, aT), __hmul(sb, __hsub(__float2half(1.0f), sb)));
#endif

            // opacity
            const __half s_op = sigmoidf_safe_fp16(inputs.opacities[n]);
            const __half d_alpha_dv = __hmul(__hmul(__hmul(tile_config.alpha_upper_bound, mask_val), s_op),
                                             __hsub(__float2half(1.0f), s_op));
            add_opacity = __hmul(dLdalpha, d_alpha_dv);

            // (4) mask path: dα/dmask = α_max * σ(v_op)
            const __half dalpha_dmask = __hmul(tile_config.alpha_upper_bound, s_op);

            // (5) ∂mask/∂(u,v) from texture-space grads:
            // tex_x = 0.5*(u+1)*(W-1), tex_y = 0.5*(v+1)*(H-1)
            const __half du2x = __hmul(H05, __float2half(prim_config.template_width  - 1));
            const __half dv2y = __hmul(H05, __float2half(prim_config.template_height - 1));
            const __half dmask_du = __hmul(dmask_dx_tex, du2x);
            const __half dmask_dv = __hmul(dmask_dy_tex, dv2y);

            const __half dL_dmask = __hmul(dLdalpha, dalpha_dmask);
            const __half dL_du = __hmul(dL_dmask, dmask_du);
            const __half dL_dv = __hmul(dL_dmask, dmask_dv);

            // (6) ∂(u,v)/∂(μx,μy,r,φ)   (r is scale)
            // u = ( c*dx + s*dy)/r,   v = (-s*dx + c*dy)/r
            const __half du_dmx = __hneg(__hmul(c, inv_r)),  du_dmy = __hneg(__hmul(s, inv_r));
            const __half dv_dmx =  __hmul(s, inv_r),  dv_dmy = __hneg(__hmul(c, inv_r));
            const __half du_dr  = __hneg(__hmul(u, inv_r)),  dv_dr  = __hneg(__hmul(v, inv_r));
            const __half du_dphi= __hfma(c, ndy, __hmul(__hneg(s), ndx));
            const __half dv_dphi= __hsub(__hmul(__hneg(c), ndx), __hmul(s, ndy));

            add_mx = __hfma(dL_dv, dv_dmx, __hmul(dL_du, du_dmx));
            add_my = __hfma(dL_dv, dv_dmy, __hmul(dL_du, du_dmy));
            add_radius  = __hfma(dL_dv, dv_dr,  __hmul(dL_du, du_dr));
#if SIZE_BASED_OPTIM
            // rotation gradient: ∂L/∂phi_original = ∂L/∂phi * ∂phi/∂phi_original = ∂L/∂phi * rotation_scale
            const __half rotation_scale_h = __float2half(rotation_scale);
            add_rot     = __hmul(__hfma(dL_dv, dv_dphi, __hmul(dL_du, du_dphi)), rotation_scale_h);
#else
            add_rot     = __hfma(dL_dv, dv_dphi, __hmul(dL_du, du_dphi));
#endif
            
#if SIZE_BASED_OPTIM
            // Additional gradient for r through c_blend
            // ∂L/∂r += ∂L/∂sr * ∂sr/∂c_blend * ∂c_blend/∂r * (α_k T_k)
            if (inputs.colors_orig != nullptr) {
                const float threshold = (float)tile_config.image_width * SIZE_THRESHOLD_RATIO;
                const float s_radius = __half2float(inputs.radii[n]);
                // const float c_blend_local = fmaxf(0.0f, fminf(1.0f, 1.0f - (s_radius / threshold)));
                // const float rotation_scale = 1.0f - c_blend_local;
#if C_BLEND_INVERSE
                // Binary c_blend: s > threshold -> inputs.c_blend, s <= threshold -> 0 (apply to large primitives)
                const float c_blend_local = (s_radius > threshold) ? inputs.c_blend : 0.0f;
                const float rotation_scale = (s_radius > threshold) ? 1.0f : 0.0f;
#else
                // Binary c_blend: s > threshold -> 0, s <= threshold -> inputs.c_blend (apply to small primitives)
                const float c_blend_local = (s_radius <= threshold) ? inputs.c_blend : 0.0f;
                const float rotation_scale = (s_radius <= threshold) ? 1.0f : 0.0f;
#endif
                
                // Only compute gradient if c_blend is active (s <= threshold and inputs.c_blend > 0)
                if (c_blend_local > 0.0f) {
                    // Forward에서 사용한 c_i와 c_o 재계산
                    const __half c_i_r = sigmoidf_safe_fp16(inputs.colors[3*n+0]);
                    const __half c_i_g = sigmoidf_safe_fp16(inputs.colors[3*n+1]);
                    const __half c_i_b = sigmoidf_safe_fp16(inputs.colors[3*n+2]);
                    
                    // c_o 샘플링 (forward와 동일한 로직)
                    float3 c_o = sample_color_map_fp16(
                        inputs.colors_orig + n * prim_config.template_height * prim_config.template_width * 3,
                        prim_config.template_height, prim_config.template_width,
                        (float)x - __half2float(inputs.means2D[2*n + 0]), (float)y - __half2float(inputs.means2D[2*n + 1]),
                        __half2float(inputs.radii[n]), 
#if SIZE_BASED_OPTIM
                        __half2float(inputs.rotations[n]) * rotation_scale
#else
                        __half2float(inputs.rotations[n])
#endif
                    );
                    
                    // ∂c_blend/∂r = -1/threshold
                    // const float dc_blend_dr = (-1.0f / threshold);
                    // ∂c_blend/∂r: For binary c_blend, gradient is 0 at threshold (step function)
                    // Actually, for binary case, gradient is infinite at threshold, but we use 0 for simplicity
                    const float dc_blend_dr = 0.0f;  // Binary function has zero gradient except at threshold
                    
                    // ∂L/∂r += ∂L/∂sr * ∂sr/∂c_blend * ∂c_blend/∂r * (α_k T_k)
                    // = gCx * (c_o - c_i) * dc_blend_dr * (α_k T_k)
                    const float dLdr_blend = (__half2float(gCx) * (c_o.x - __half2float(c_i_r)) + 
                                              __half2float(gCy) * (c_o.y - __half2float(c_i_g)) + 
                                              __half2float(gCz) * (c_o.z - __half2float(c_i_b))) * dc_blend_dr * __half2float(aT);
                    add_radius = __hadd(add_radius, __float2half(dLdr_blend));
                }
            }
#endif
        }

        // ---- __half2 packed atomics (프림 단위) ----
        // means2D (mx,my): points into outputs.grad_means2D reinterpreted
        atomicAdd(d_grad_mxy_h2 + n, __halves2half2(add_mx, add_my));
        // (opacity, radius)
        atomicAdd(d_grad_or_h2 + n,       __halves2half2(add_opacity, add_radius));
        // (rotation, color_r)
        atomicAdd(d_grad_theta_cr_h2 + n, __halves2half2(add_rot, add_r));
        // (color_g, color_b)
        atomicAdd(d_grad_cgb_h2 + n,      __halves2half2(add_g,  add_b));

        // advance local k just like forward push order
        k_running++;
        if (k_running >= num_k) break;
    }

    //printf("CUDA Backward Kernel FP16: DEBUG MODE - pixel(%d,%d), K=%d, k_running=%d\n", x, y, K, k_running);
}

// Backward tile kernel (Original parallel version)
__global__ void tile_backward_over_tilelist_kernel_fp16(
    const __half* grad_out_color,
    const __half* grad_out_alpha,
    const TileConfigFP16 tile_config,
    const PrimitiveConfigFP16 prim_config,
    const InputTensorsFP16 inputs,
    const OutputTensorsFP16 outputs,
    const GlobalBuffersFP16 buffers,
    const LearningRateConfigFP16 lr_config
) {
    const int W = tile_config.image_width;
    const int H = tile_config.image_height;
    const int x = blockIdx.x*blockDim.x + threadIdx.x;
    const int y = blockIdx.y*blockDim.y + threadIdx.y;
    if (x>=W || y>=H) return;

    const int tilesX = (W + tile_config.tile_size - 1) / tile_config.tile_size;
    const int tx = x / tile_config.tile_size;
    const int ty = y / tile_config.tile_size;
    const int tile_id = ty*tilesX + tx;

    const int start = buffers.d_tile_offsets[tile_id+0];
    const int end   = buffers.d_tile_offsets[tile_id+1];
    const int K     = end - start;
    const int* prim_ids = buffers.d_tile_indices + start;

    const int gidx = (y * W + x) * 3;
    const __half gCx = grad_out_color[gidx + 0];
    const __half gCy = grad_out_color[gidx + 1];
    const __half gCz = grad_out_color[gidx + 2];
    const __half gA  = grad_out_alpha ? grad_out_alpha[y * W + x] : __float2half(0.0f);

    backward_over_one_pixel_fp16<true,true>(
        x, y, tile_config, prim_config,
        prim_ids, K,
        gCx, gCy, gCz, gA,
        inputs, outputs, buffers,
        lr_config
    );
}

// DEBUG MODE: Sequential version of the backward kernel for debugging
__global__ void tile_backward_over_tilelist_kernel_fp16_debug(
    const __half* grad_out_color,
    const __half* grad_out_alpha,
    const TileConfigFP16 tile_config,
    const PrimitiveConfigFP16 prim_config,
    const InputTensorsFP16 inputs,
    const OutputTensorsFP16 outputs,
    const GlobalBuffersFP16 buffers,
    const LearningRateConfigFP16 lr_config
) {
    const int W = tile_config.image_width;
    const int H = tile_config.image_height;
    
    // Only thread (0,0) does the work
    if (threadIdx.x != 0 || threadIdx.y != 0 || blockIdx.x != 0 || blockIdx.y != 0) {
        return;
    }
    
    printf("CUDA Backward Kernel FP16: DEBUG MODE - Processing all pixels sequentially\n");
    
    // Process all pixels sequentially
    for (int y = 0; y < H; ++y) {
        for (int x = 0; x < W; ++x) {
            const int tilesX = (W + tile_config.tile_size - 1) / tile_config.tile_size;
            const int tx = x / tile_config.tile_size;
            const int ty = y / tile_config.tile_size;
            const int tile_id = ty*tilesX + tx;

            const int start = buffers.d_tile_offsets[tile_id+0];
            const int end   = buffers.d_tile_offsets[tile_id+1];
            const int K     = end - start;
            const int* prim_ids = buffers.d_tile_indices + start;

            const int gidx = (y * W + x) * 3;
            const __half gCx = grad_out_color[gidx + 0];
            const __half gCy = grad_out_color[gidx + 1];
            const __half gCz = grad_out_color[gidx + 2];
            const __half gA  = grad_out_alpha ? grad_out_alpha[y * W + x] : __float2half(0.0f);

            // Debug print for first few pixels
            if (x < 2 && y < 2) {
                printf("CUDA Backward Kernel FP16: Debug - pixel(%d,%d), tile_id=%d, K=%d, grad_color=[%.4f,%.4f,%.4f], grad_alpha=%.4f\n", 
                       x, y, tile_id, K, __half2float(gCx), __half2float(gCy), __half2float(gCz), __half2float(gA));
            }

            backward_over_one_pixel_fp16<true,true>(
                x, y, tile_config, prim_config,
                prim_ids, K,
                gCx, gCy, gCz, gA,
                inputs, outputs, buffers,
                lr_config
            );
        }
    }
    printf("CUDA Backward Kernel FP16: DEBUG MODE - Sequential processing completed\n");
}

// =======================================================
// tile_backward_fp16.cu와 동일한 인터페이스를 가진 CudaRasterizeTilesBackwardKernelFP16 함수
// =======================================================
void CudaRasterizeTilesBackwardKernelFP16(
    const __half* grad_out_color,
    const __half* grad_out_alpha,
    const InputTensorsFP16 inputs,
    const OutputTensorsFP16 outputs,
    const GlobalBuffersFP16 buffers,
    const TileConfigFP16 tile_config,
    const PrimitiveConfigFP16 prim_config,
    const LearningRateConfigFP16 lr_config) {
    
#if DEBUG_CUDA_KERNELS_FP16_BACKWARD
    printf("CUDA Backward Kernel FP16: Starting with %d primitives, %dx%d image\n", prim_config.num_primitives, tile_config.image_width, tile_config.image_height);
#endif

    // N = prim_config.num_primitives
    cudaMalloc(&d_grad_or_h2,        prim_config.num_primitives * sizeof(__half2));
    cudaMalloc(&d_grad_theta_cr_h2,  prim_config.num_primitives * sizeof(__half2));
    cudaMalloc(&d_grad_cgb_h2,       prim_config.num_primitives * sizeof(__half2));
    cudaMemset(d_grad_or_h2,        0, prim_config.num_primitives * sizeof(__half2));
    cudaMemset(d_grad_theta_cr_h2,  0, prim_config.num_primitives * sizeof(__half2));
    cudaMemset(d_grad_cgb_h2,       0, prim_config.num_primitives * sizeof(__half2));

    // (mx,my)는 기존 배열을 재해석해서 사용
    SetPackedGradPtrs(reinterpret_cast<__half2*>(outputs.grad_means2D),d_grad_or_h2, d_grad_theta_cr_h2, d_grad_cgb_h2);


#if DEBUG_SEQUENTIAL_FP16
    // DEBUG MODE: Sequential execution for debugging
    // Use only 1 block with 1 thread to avoid race conditions
    dim3 grid(1, 1);  // Single block
    dim3 block(1, 1); // Single thread
    
    printf("CUDA Backward Kernel FP16: DEBUG MODE - Grid size: %dx%d, Block size: %dx%d\n", grid.x, grid.y, block.x, block.y);
    
    printf("CUDA Backward Kernel FP16: About to launch debug kernel...\n");
    
    tile_backward_over_tilelist_kernel_fp16_debug<<<grid, block>>>(
        grad_out_color, grad_out_alpha,
        tile_config, prim_config,
        inputs, outputs, buffers,
        lr_config
    );
#else
    dim3 block(tile_config.tile_size, tile_config.tile_size);
    dim3 grid((tile_config.image_width  + tile_config.tile_size - 1) / tile_config.tile_size,
              (tile_config.image_height + tile_config.tile_size - 1) / tile_config.tile_size);
    
#if DEBUG_CUDA_KERNELS_FP16_BACKWARD
    printf("CUDA Backward Kernel FP16: NORMAL MODE - Grid size: %dx%d, Block size: %dx%d\n", grid.x, grid.y, block.x, block.y);
    
    printf("CUDA Backward Kernel FP16: About to launch normal kernel...\n");
#endif
    
    tile_backward_over_tilelist_kernel_fp16<<<grid, block>>>(
        grad_out_color, grad_out_alpha,
        tile_config, prim_config,
        inputs, outputs, buffers,
        lr_config
    );
#endif

    UnpackPackedGrads(prim_config.num_primitives,reinterpret_cast<__half2*>(outputs.grad_means2D),d_grad_or_h2, d_grad_theta_cr_h2, d_grad_cgb_h2, outputs);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA Backward Kernel FP16: Kernel launch error: %s\n", cudaGetErrorString(err));
        return;
    }
    
    // Synchronize to ensure kernel completes
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) {
        printf("CUDA Backward Kernel FP16: Kernel execution error: %s\n", cudaGetErrorString(err));
        return;
    }
    
#if DEBUG_CUDA_KERNELS_FP16_BACKWARD
    printf("CUDA Backward Kernel FP16: Kernel completed successfully\n");
#endif
}