// ===================================================================
// Element-wise and memory-layout kernels for GpuTensor -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), 0.0f, 1.0f,
// expf(), sqrtf(), fabsf(), fmaxf(), etc.
//
// All kernels use 1-D grid-stride pattern with 256 threads per block.
// ===================================================================

// -- Binary element-wise: add, sub, mul, div, min, max ----------------

extern "C" __global__
void elementwise_binary(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ out,
    int n,
    int op
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float va = a[i];
        float vb = b[i];
        float r;
        float zero = 0.0f;
        float eps = (1e-10f);
        switch (op) {
            case 0: r = va + vb; break;  // add
            case 1: r = va - vb; break;  // sub
            case 2: r = va * vb; break;  // mul
            case 3: r = (fabsf(vb) < eps) ? zero : va / vb; break;  // div
            case 4: r = fminf(va, vb); break;  // min
            case 5: r = fmaxf(va, vb); break;  // max
            default: r = zero; break;
        }
        out[i] = r;
    }
}

// -- Unary element-wise: powf..affine, sin, cos, sqrt, silu, elu, recip --

extern "C" __global__
void elementwise_unary(
    const float* __restrict__ x,
    float* __restrict__ out,
    int n,
    int op,
    float param1,
    float param2
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float v = x[i];
        float r;
        float zero = 0.0f;
        float one = 1.0f;
        float p1 = (param1);
        float p2 = (param2);
        switch (op) {
            case 0: r = powf(v, param1); break;           // powf
            case 1: r = v * v; break;                            // sqr
            case 2: r = floorf(v); break;                    // floor
            case 3: r = fmaxf(v, zero); break;               // relu
            case 4: r = fminf(fmaxf(v, p1), p2); break; // clamp
            case 5: r = fabsf(v); break;                     // abs
            case 6: r = zero - v; break;                         // neg
            case 7: r = expf(v); break;                      // exp
            case 8: r = logf(v); break;                      // log
            case 9: r = (v <= p1) ? one : zero; break;           // le (compare)
            case 10: r = v * p1 + p2; break;                     // affine (mul + add)
            case 11: r = sinf(v); break;                     // sin
            case 12: r = cosf(v); break;                     // cos
            case 13: r = sqrtf(v); break;                    // sqrt
            case 14: {                                            // silu
                float s = one / (one + expf(zero - v));
                r = v * s;
                break;
            }
            case 15: {                                            // elu
                r = (v >= zero) ? v : p1 * (expf(v) - one);
                break;
            }
            case 16: {                                            // recip
                float eps = (1e-10f);
                float big = (1e10f);
                r = (fabsf(v) < eps) ? big : one / v;
                break;
            }
            case 17: {                                            // sigmoid
                r = one / (one + expf(zero - v));
                break;
            }
            case 18: r = tanhf(v); break;                    // tanh
            default: r = v; break;
        }
        out[i] = r;
    }
}

// -- Broadcast scalar binary: tensor op scalar or scalar op tensor ----

extern "C" __global__
void broadcast_scalar_binary(
    const float* __restrict__ tensor,
    const float* __restrict__ scalar_ptr,
    float* __restrict__ out,
    int n,
    int op,
    int scalar_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float s = scalar_ptr[0];
        float t = tensor[i];
        float lhs = scalar_is_lhs ? s : t;
        float rhs = scalar_is_lhs ? t : s;
        float r;
        float zero = 0.0f;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = zero; break;
        }
        out[i] = r;
    }
}

// -- Broadcast row binary: [M,N] op [1,N] or [1,N] op [M,N] ---------

extern "C" __global__
void broadcast_row_binary(
    const float* __restrict__ matrix,
    const float* __restrict__ row,
    float* __restrict__ out,
    int total,
    int cols,
    int op,
    int row_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int col = i % cols;
        float m = matrix[i];
        float rv = row[col];
        float lhs = row_is_lhs ? rv : m;
        float rhs = row_is_lhs ? m  : rv;
        float r;
        float zero = 0.0f;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = zero; break;
        }
        out[i] = r;
    }
}

// -- Broadcast column binary: [M,N] op [M,1] or [M,1] op [M,N] ------

extern "C" __global__
void broadcast_col_binary(
    const float* __restrict__ matrix,
    const float* __restrict__ col_vec,
    float* __restrict__ out,
    int total,
    int cols,
    int op,
    int col_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int row_idx = i / cols;
        float m = matrix[i];
        float cv = col_vec[row_idx];
        float lhs = col_is_lhs ? cv : m;
        float rhs = col_is_lhs ? m  : cv;
        float r;
        float zero = 0.0f;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = zero; break;
        }
        out[i] = r;
    }
}

// -- Expand (general N-D broadcast) -----------------------------------

extern "C" __global__
void expand_broadcast(
    const float* __restrict__ src,
    const unsigned int* __restrict__ src_strides,
    const unsigned int* __restrict__ src_shape,
    const unsigned int* __restrict__ tgt_strides,
    float* __restrict__ out,
    int total,
    int ndim
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int src_idx = 0;
        int remaining = i;
        for (int d = 0; d < ndim; d++) {
            int tgt_stride = tgt_strides[d];
            int coord = remaining / tgt_stride;
            remaining = remaining % tgt_stride;
            int src_dim = src_shape[d];
            int src_coord = (src_dim == 1) ? 0 : coord;
            src_idx += src_coord * src_strides[d];
        }
        out[i] = src[src_idx];
    }
}

// -- Transpose 2-D ----------------------------------------------------

extern "C" __global__
void transpose_2d(
    const float* __restrict__ src,
    float* __restrict__ dst,
    int total,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int r = i / cols;
        int c = i % cols;
        dst[c * rows + r] = src[i];
    }
}

// -- Gather (index_select along a dimension) --------------------------

extern "C" __global__
void gather_select(
    const float* __restrict__ src,
    const unsigned int* __restrict__ indices,
    float* __restrict__ out,
    int total,
    int num_indices,
    int inner_size,
    int src_dim_stride
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int out_idx_inner = i % inner_size;
        int tmp = i / inner_size;
        int idx_pos = tmp % num_indices;
        int outer = tmp / num_indices;

        unsigned int src_idx = indices[idx_pos];
        int src_flat = outer * src_dim_stride + (int)src_idx * inner_size + out_idx_inner;
        out[i] = src[src_flat];
    }
}

// -- Per-row gather: out[i] = src[i * cols + clamp(indices[i], 0, cols-1)] --

extern "C" __global__
void gather_rows(
    const float* __restrict__ src,
    const int* __restrict__ indices,
    float* __restrict__ out,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < rows) {
        int col = indices[i];
        col = col < 0 ? 0 : (col >= cols ? cols - 1 : col);
        out[i] = src[i * cols + col];
    }
}

// -- Per-row gather with unsigned int indices -------------------------

extern "C" __global__
void gather_rows_u32(
    const float* __restrict__ src,
    const unsigned int* __restrict__ indices,
    float* __restrict__ out,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < rows) {
        int col = (int)indices[i];
        col = col >= cols ? cols - 1 : col;
        out[i] = src[i * cols + col];
    }
}
