// ===================================================================
// Fused reduction kernels: stats, argmax, sum, per-column sum.
// f32-native storage and accumulation.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, atomicAdd(), __shfl_down_sync(), etc.
//
// All kernels use shared-memory tree reduction with warp shuffle.
// ===================================================================

// -- Helper: warp-level reductions (f32 native) -----------------------

__device__ __forceinline__ float warp_reduce_sum_(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = val + __shfl_down_sync(0xFFFFFFFF, val, offset);
    return val;
}

__device__ __forceinline__ float warp_reduce_min_(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fminf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
    return val;
}

__device__ __forceinline__ float warp_reduce_max_(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
    return val;
}

// -- Fused stats: min, max, sum, sum_sq, count -------------------------
// result[5] as f32: [min, max, sum, sum_sq, count]
// Initialized by host: min=+INF, max=-INF, sum/sum_sq/count=0

extern "C" __global__
void fused_stats_reduce(const float* __restrict__ data,
                        float* __restrict__ result,
                        int n) {
    extern __shared__ float shmem_f32[];
    int B = blockDim.x;

    float local_min = (3.402823466e+38f);
    float local_max = (-3.402823466e+38f);
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    float local_count = 0.0f;
    float one = 1.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        float v = data[i];
        local_min = fminf(local_min, v);
        local_max = fmaxf(local_max, v);
        local_sum = local_sum + v;
        local_sum_sq = local_sum_sq + v * v;
        local_count = local_count + one;
    }

    shmem_f32[threadIdx.x]         = local_min;
    shmem_f32[B + threadIdx.x]     = local_max;
    shmem_f32[2*B + threadIdx.x]   = local_sum;
    shmem_f32[3*B + threadIdx.x]   = local_sum_sq;
    shmem_f32[4*B + threadIdx.x]   = local_count;
    __syncthreads();

    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem_f32[threadIdx.x]         = fminf(shmem_f32[threadIdx.x], shmem_f32[threadIdx.x + s]);
            shmem_f32[B + threadIdx.x]     = fmaxf(shmem_f32[B + threadIdx.x], shmem_f32[B + threadIdx.x + s]);
            shmem_f32[2*B + threadIdx.x]   = shmem_f32[2*B + threadIdx.x] + shmem_f32[2*B + threadIdx.x + s];
            shmem_f32[3*B + threadIdx.x]   = shmem_f32[3*B + threadIdx.x] + shmem_f32[3*B + threadIdx.x + s];
            shmem_f32[4*B + threadIdx.x]   = shmem_f32[4*B + threadIdx.x] + shmem_f32[4*B + threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        float w_min = fminf(shmem_f32[threadIdx.x], shmem_f32[threadIdx.x + 32]);
        float w_max = fmaxf(shmem_f32[B + threadIdx.x], shmem_f32[B + threadIdx.x + 32]);
        float w_sum = shmem_f32[2*B + threadIdx.x] + shmem_f32[2*B + threadIdx.x + 32];
        float w_ssq = shmem_f32[3*B + threadIdx.x] + shmem_f32[3*B + threadIdx.x + 32];
        float w_cnt = shmem_f32[4*B + threadIdx.x] + shmem_f32[4*B + threadIdx.x + 32];

        w_min = warp_reduce_min_(w_min);
        w_max = warp_reduce_max_(w_max);
        w_sum = warp_reduce_sum_(w_sum);
        w_ssq = warp_reduce_sum_(w_ssq);
        w_cnt = warp_reduce_sum_(w_cnt);

        if (threadIdx.x == 0) {
            // atomicMin for f32 via CAS loop (unsigned int = 4 bytes = sizeof(float))
            {
                unsigned int* addr_ui = (unsigned int*)&result[0];
                unsigned int old_bits = *addr_ui;
                unsigned int new_bits = __float_as_uint(w_min);
                while (__uint_as_float(old_bits) > w_min) {
                    unsigned int assumed = old_bits;
                    old_bits = atomicCAS(addr_ui, assumed, new_bits);
                    if (old_bits == assumed) break;
                    float cur = __uint_as_float(old_bits);
                    float new_val = fminf(cur, w_min);
                    new_bits = __float_as_uint(new_val);
                }
            }
            // atomicMax for f32 via CAS loop
            {
                unsigned int* addr_ui = (unsigned int*)&result[1];
                unsigned int old_bits = *addr_ui;
                unsigned int new_bits = __float_as_uint(w_max);
                while (__uint_as_float(old_bits) < w_max) {
                    unsigned int assumed = old_bits;
                    old_bits = atomicCAS(addr_ui, assumed, new_bits);
                    if (old_bits == assumed) break;
                    float cur = __uint_as_float(old_bits);
                    float new_val = fmaxf(cur, w_max);
                    new_bits = __float_as_uint(new_val);
                }
            }
            atomicAdd(&result[2], w_sum);
            atomicAdd(&result[3], w_ssq);
            atomicAdd(&result[4], w_cnt);
        }
    }
}

// -- Argmax (flat buffer) ---------------------------------------------
// NOTE: argmax needs int index tracking alongside f32 values.
// The index is stored via __int_as_float/__float_as_int bit tricks in
// f32 shared memory.

extern "C" __global__
void argmax_flat(const float* __restrict__ data,
                 float* __restrict__ result,
                 int n) {
    // We need to track both values and indices through reduction.
    extern __shared__ char shmem_raw[];
    int B = blockDim.x;
    // Layout: [B f32 values][B int indices]
    float* shmem_val = (float*)shmem_raw;
    int* shmem_idx = (int*)(shmem_raw + B * sizeof(float));

    float local_val = (-3.402823466e+38f);
    int local_idx = 0;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        float v = data[i];
        if (v > local_val) {
            local_val = v;
            local_idx = i;
        }
    }

    shmem_val[threadIdx.x] = local_val;
    shmem_idx[threadIdx.x] = local_idx;
    __syncthreads();

    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            if (shmem_val[threadIdx.x + s] > shmem_val[threadIdx.x]) {
                shmem_val[threadIdx.x] = shmem_val[threadIdx.x + s];
                shmem_idx[threadIdx.x] = shmem_idx[threadIdx.x + s];
            }
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        float w_val = shmem_val[threadIdx.x];
        int w_idx = shmem_idx[threadIdx.x];
        float other_val32 = shmem_val[threadIdx.x + 32];
        int other_idx32 = shmem_idx[threadIdx.x + 32];
        if (other_val32 > w_val) {
            w_val = other_val32;
            w_idx = other_idx32;
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
            int other_idx_s = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
            if (other_val > w_val) {
                w_val = other_val;
                w_idx = other_idx_s;
            }
        }

        if (threadIdx.x == 0) {
            // CAS loop to merge block results into global output
            // result[0] = best val (as f32), result[1] = best idx (as f32)
            unsigned int* val_addr = (unsigned int*)&result[0];
            unsigned int old_bits = *val_addr;
            while (__uint_as_float(old_bits) < w_val) {
                unsigned int assumed = old_bits;
                old_bits = atomicCAS(val_addr, assumed, __float_as_uint(w_val));
                if (old_bits == assumed) {
                    result[1] = (float)w_idx;
                    break;
                }
            }
        }
    }
}

// -- Argmax per row for 2D data ----------------------------------------

extern "C" __global__
void argmax_rows(const float* __restrict__ data,
                 unsigned int* __restrict__ result,
                 int rows,
                 int cols) {
    extern __shared__ char shmem_raw2[];
    int B = blockDim.x;
    int row = blockIdx.x;
    if (row >= rows) return;

    // Layout: [B f32 values][B int indices]
    float* shmem_val = (float*)shmem_raw2;
    int* shmem_idx = (int*)(shmem_raw2 + B * sizeof(float));

    const float* row_data = data + row * cols;

    float local_val = (-3.402823466e+38f);
    int local_idx = 0;

    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
        float v = row_data[c];
        if (v > local_val) {
            local_val = v;
            local_idx = c;
        }
    }

    shmem_val[threadIdx.x] = local_val;
    shmem_idx[threadIdx.x] = local_idx;
    __syncthreads();

    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            if (shmem_val[threadIdx.x + s] > shmem_val[threadIdx.x]) {
                shmem_val[threadIdx.x] = shmem_val[threadIdx.x + s];
                shmem_idx[threadIdx.x] = shmem_idx[threadIdx.x + s];
            }
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        float w_val = shmem_val[threadIdx.x];
        int w_idx = shmem_idx[threadIdx.x];
        if (B >= 64) {
            float other_val32 = shmem_val[threadIdx.x + 32];
            int other_idx32 = shmem_idx[threadIdx.x + 32];
            if (other_val32 > w_val) {
                w_val = other_val32;
                w_idx = other_idx32;
            }
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
            int other_idx_s = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
            if (other_val > w_val) {
                w_val = other_val;
                w_idx = other_idx_s;
            }
        }

        if (threadIdx.x == 0) {
            result[row] = (unsigned int)w_idx;
        }
    }
}

// -- Sum reduction ----------------------------------------------------

extern "C" __global__
void sum_reduce(const float* __restrict__ data,
                float* __restrict__ result,
                int n) {
    extern __shared__ float shmem_sum[];

    float local_sum = 0.0f;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        local_sum = local_sum + data[i];
    }

    shmem_sum[threadIdx.x] = local_sum;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem_sum[threadIdx.x] = shmem_sum[threadIdx.x] + shmem_sum[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        float w_sum = shmem_sum[threadIdx.x] + shmem_sum[threadIdx.x + 32];
        w_sum = warp_reduce_sum_(w_sum);
        if (threadIdx.x == 0) {
            atomicAdd(result, w_sum);
        }
    }
}

// -- Per-column sum for 2D data ----------------------------------------

extern "C" __global__
void col_sum_reduce(const float* __restrict__ data,
                    float* __restrict__ result,
                    int rows,
                    int cols) {
    extern __shared__ float shmem_col[];
    int col = blockIdx.x;
    if (col >= cols) return;

    float local_sum = 0.0f;

    for (int r = threadIdx.x; r < rows; r += blockDim.x) {
        local_sum = local_sum + data[r * cols + col];
    }

    shmem_col[threadIdx.x] = local_sum;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem_col[threadIdx.x] = shmem_col[threadIdx.x] + shmem_col[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        float w_sum = shmem_col[threadIdx.x];
        if (blockDim.x >= 64)
            w_sum = w_sum + shmem_col[threadIdx.x + 32];
        w_sum = warp_reduce_sum_(w_sum);
        if (threadIdx.x == 0) {
            result[col] = w_sum;
        }
    }
}
