#include <cuda_bf16.h>
#include <math.h>
#include "../common.cuh"

__global__ void depthwise_causal_conv1d_silu_bf16_kernel(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ w,
    __nv_bfloat16* __restrict__ out,
    int seq_len,
    int channels,
    int kernel
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = seq_len * channels;
    if (idx >= total) return;

    int c = idx % channels;
    int t = idx / channels;
    float acc = 0.0f;
    for (int k = 0; k < kernel; ++k) {
        int src_t = t - (kernel - 1 - k);
        if (src_t >= 0) {
            float xv = __bfloat162float(x[src_t * channels + c]);
            float wv = __bfloat162float(w[c * kernel + k]);
            acc += xv * wv;
        }
    }
    float y = acc / (1.0f + expf(-acc));
    out[idx] = __float2bfloat16(y);
}

__global__ void deltanet_ar_bf16_kernel(
    const __nv_bfloat16* __restrict__ q,
    const __nv_bfloat16* __restrict__ k,
    const __nv_bfloat16* __restrict__ v,
    const __nv_bfloat16* __restrict__ beta,
    const __nv_bfloat16* __restrict__ alpha_softplus,
    const __nv_bfloat16* __restrict__ a_log,
    __nv_bfloat16* __restrict__ state,
    __nv_bfloat16* __restrict__ out,
    int key_heads,
    int value_heads,
    int head_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = value_heads * head_dim;
    if (idx >= total) return;

    int h = idx / head_dim;
    int j = idx % head_dim;
    int key_h = h * key_heads / value_heads;

    float beta_h = __bfloat162float(beta[h]);
    float gate = expf(__bfloat162float(alpha_softplus[h]) * __bfloat162float(a_log[h]));

    float sk = 0.0f;
    for (int i = 0; i < head_dim; ++i) {
        float s = __bfloat162float(state[(h * head_dim + i) * head_dim + j]);
        float ki = __bfloat162float(k[key_h * head_dim + i]);
        s *= gate;
        state[(h * head_dim + i) * head_dim + j] = __float2bfloat16(s);
        sk += s * ki;
    }

    float vj = __bfloat162float(v[h * head_dim + j]);
    float d = (vj - sk) * beta_h;

    float out_j = 0.0f;
    for (int i = 0; i < head_dim; ++i) {
        float ki = __bfloat162float(k[key_h * head_dim + i]);
        float qi = __bfloat162float(q[key_h * head_dim + i]) / sqrtf((float)head_dim);
        int sidx = (h * head_dim + i) * head_dim + j;
        float s = __bfloat162float(state[sidx]) + ki * d;
        state[sidx] = __float2bfloat16(s);
        out_j += s * qi;
    }
    out[h * head_dim + j] = __float2bfloat16(out_j);
}

// Stateful depthwise causal convolution for hybrid recurrent attention.
// `state` stores the preceding `kernel - 1` inputs in chronological order.
// One thread owns a channel, so the state update is race-free for any seq_len.
__global__ void depthwise_causal_conv1d_stateful_silu_bf16_kernel(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ w,
    __nv_bfloat16* __restrict__ state,
    __nv_bfloat16* __restrict__ out,
    int seq_len,
    int channels,
    int kernel
) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= channels) return;

    const int history = kernel - 1;
    for (int t = 0; t < seq_len; ++t) {
        float acc = 0.0f;
        for (int k = 0; k < kernel; ++k) {
            int src_t = t - history + k;
            float xv;
            if (src_t < 0) {
                xv = __bfloat162float(state[c * history + history + src_t]);
            } else {
                xv = __bfloat162float(x[(long long)src_t * channels + c]);
            }
            acc += xv * __bfloat162float(w[c * kernel + k]);
        }
        out[(long long)t * channels + c] =
            __float2bfloat16(acc / (1.0f + expf(-acc)));
    }

    // Keep the last `history` raw projection values for the next call.
    for (int h = 0; h < history; ++h) {
        int src_t = seq_len - history + h;
        __nv_bfloat16 value;
        if (src_t < 0) {
            value = state[c * history + history + src_t];
        } else {
            value = x[(long long)src_t * channels + c];
        }
        state[c * history + h] = value;
    }
}

// L2-normalize each row using FP32 accumulation. Qwen3.5/3.6 applies this to
// convolved Q and K before the delta rule (this is not RMSNorm).
__global__ void l2_normalize_rows_bf16_kernel(
    const __nv_bfloat16* __restrict__ x,
    __nv_bfloat16* __restrict__ out,
    int rows,
    int cols,
    float eps
) {
    int row = blockIdx.x;
    if (row >= rows) return;

    float local = 0.0f;
    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
        float v = __bfloat162float(x[(long long)row * cols + i]);
        local += v * v;
    }
    __shared__ float sums[256];
    sums[threadIdx.x] = local;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) sums[threadIdx.x] += sums[threadIdx.x + stride];
        __syncthreads();
    }
    // Match torch.nn.functional.normalize / llama.cpp: clamp the norm by eps,
    // rather than adding eps to every non-zero norm.
    float inv = rsqrtf(fmaxf(sums[0], eps * eps));
    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
        out[(long long)row * cols + i] =
            __float2bfloat16(__bfloat162float(x[(long long)row * cols + i]) * inv);
    }
}

// Recurrent Gated DeltaNet over one or more tokens. HF stores V heads grouped
// by K head, so V head h uses K/Q head floor(h / (value_heads/key_heads)).
// The recurrent matrix is FP32 as required by `mamba_ssm_dtype=float32`.
__global__ void deltanet_recurrent_bf16_f32_kernel(
    const __nv_bfloat16* __restrict__ q,
    const __nv_bfloat16* __restrict__ k,
    const __nv_bfloat16* __restrict__ v,
    const __nv_bfloat16* __restrict__ beta,
    const __nv_bfloat16* __restrict__ alpha_softplus,
    const __nv_bfloat16* __restrict__ a_log,
    float* __restrict__ state,
    __nv_bfloat16* __restrict__ out,
    int seq_len,
    int key_heads,
    int value_heads,
    int head_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = value_heads * head_dim;
    if (idx >= total) return;

    int h = idx / head_dim;
    int j = idx % head_dim;
    int key_h = h * key_heads / value_heads;
    float a = __bfloat162float(a_log[h]);
    float neg_a = -expf(a);
    float q_scale = rsqrtf((float)head_dim);

    for (int t = 0; t < seq_len; ++t) {
        const __nv_bfloat16* q_t = q + ((long long)t * key_heads + key_h) * head_dim;
        const __nv_bfloat16* k_t = k + ((long long)t * key_heads + key_h) * head_dim;
        const __nv_bfloat16* v_t = v + ((long long)t * value_heads + h) * head_dim;
        float b = __bfloat162float(beta[(long long)t * value_heads + h]);
        float alpha = __bfloat162float(alpha_softplus[(long long)t * value_heads + h]);
        float decay = expf(neg_a * alpha);

        float sk = 0.0f;
        for (int i = 0; i < head_dim; ++i) {
            long long sidx = ((long long)h * head_dim + i) * head_dim + j;
            float s = state[sidx] * decay;
            state[sidx] = s;
            sk += s * __bfloat162float(k_t[i]);
        }

        float d = (__bfloat162float(v_t[j]) - sk) * b;
        float out_j = 0.0f;
        for (int i = 0; i < head_dim; ++i) {
            long long sidx = ((long long)h * head_dim + i) * head_dim + j;
            float s = state[sidx] + __bfloat162float(k_t[i]) * d;
            state[sidx] = s;
            out_j += s * (__bfloat162float(q_t[i]) * q_scale);
        }
        out[((long long)t * value_heads + h) * head_dim + j] = __float2bfloat16(out_j);
    }
}

extern "C" {

void launch_depthwise_causal_conv1d_silu_bf16(
    const void* x,
    const void* w,
    void* out,
    int seq_len,
    int channels,
    int kernel,
    void* stream
) {
    int total = seq_len * channels;
    int block = 256;
    int grid = (total + block - 1) / block;
    depthwise_causal_conv1d_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)x,
        (const __nv_bfloat16*)w,
        (__nv_bfloat16*)out,
        seq_len,
        channels,
        kernel
    );
    CUDA_CHECK_LAST_ERROR();
}

void launch_deltanet_ar_bf16(
    const void* q,
    const void* k,
    const void* v,
    const void* beta,
    const void* alpha_softplus,
    const void* a_log,
    void* state,
    void* out,
    int key_heads,
    int value_heads,
    int head_dim,
    void* stream
) {
    int total = value_heads * head_dim;
    int block = 128;
    int grid = (total + block - 1) / block;
    deltanet_ar_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)q,
        (const __nv_bfloat16*)k,
        (const __nv_bfloat16*)v,
        (const __nv_bfloat16*)beta,
        (const __nv_bfloat16*)alpha_softplus,
        (const __nv_bfloat16*)a_log,
        (__nv_bfloat16*)state,
        (__nv_bfloat16*)out,
        key_heads,
        value_heads,
        head_dim
    );
    CUDA_CHECK_LAST_ERROR();
}

void launch_depthwise_causal_conv1d_stateful_silu_bf16(
    const void* x,
    const void* w,
    void* state,
    void* out,
    int seq_len,
    int channels,
    int kernel,
    void* stream
) {
    int block = 256;
    int grid = (channels + block - 1) / block;
    depthwise_causal_conv1d_stateful_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)x,
        (const __nv_bfloat16*)w,
        (__nv_bfloat16*)state,
        (__nv_bfloat16*)out,
        seq_len,
        channels,
        kernel
    );
    CUDA_CHECK_LAST_ERROR();
}

void launch_l2_normalize_rows_bf16(
    const void* x,
    void* out,
    int rows,
    int cols,
    float eps,
    void* stream
) {
    l2_normalize_rows_bf16_kernel<<<rows, 256, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)x,
        (__nv_bfloat16*)out,
        rows,
        cols,
        eps
    );
    CUDA_CHECK_LAST_ERROR();
}

void launch_deltanet_recurrent_bf16_f32(
    const void* q,
    const void* k,
    const void* v,
    const void* beta,
    const void* alpha_softplus,
    const void* a_log,
    void* state,
    void* out,
    int seq_len,
    int key_heads,
    int value_heads,
    int head_dim,
    void* stream
) {
    int total = value_heads * head_dim;
    int block = 128;
    int grid = (total + block - 1) / block;
    deltanet_recurrent_bf16_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)q,
        (const __nv_bfloat16*)k,
        (const __nv_bfloat16*)v,
        (const __nv_bfloat16*)beta,
        (const __nv_bfloat16*)alpha_softplus,
        (const __nv_bfloat16*)a_log,
        (float*)state,
        (__nv_bfloat16*)out,
        seq_len,
        key_heads,
        value_heads,
        head_dim
    );
    CUDA_CHECK_LAST_ERROR();
}

}