#include "neuroflow/swiglu.hpp"

#include <algorithm>
#include <cmath>
#include <cstring>

#ifdef USE_CBLAS
#include <cblas.h>
#endif

#ifdef USE_CUDA
#include "cuda_context.hpp"
#endif

namespace neuroflow {

static Tensor linear_backward_weight(const Tensor& input, const Tensor& output_grad) {
    size_t in_dim = input.shape_.back();
    size_t out_dim = output_grad.shape_.back();
    size_t batch = input.numel() / in_dim;

    Tensor grad({out_dim, in_dim}, QuantType::FP32);
    float* gp = grad.as_fp32();
    memset(gp, 0, grad.data_size_);

    const float* ip = input.as_fp32();
    const float* op = output_grad.as_fp32();

#ifdef USE_CBLAS
    cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
        static_cast<int>(out_dim), static_cast<int>(in_dim), static_cast<int>(batch),
        1.0f, op, static_cast<int>(out_dim), ip, static_cast<int>(in_dim),
        0.0f, gp, static_cast<int>(in_dim));
#else
    for (size_t b = 0; b < batch; ++b) {
        for (size_t o = 0; o < out_dim; ++o) {
            for (size_t i = 0; i < in_dim; ++i) {
                gp[o * in_dim + i] += op[b * out_dim + o] * ip[b * in_dim + i];
            }
        }
    }
#endif
    return grad;
}

static Tensor linear_backward_input(const Tensor& output_grad, const Tensor& weight) {
    size_t out_dim = output_grad.shape_.back();
    size_t batch = output_grad.numel() / out_dim;
    size_t in_dim = weight.shape_[1];

    Tensor grad({batch, in_dim}, QuantType::FP32);
    float* gp = grad.as_fp32();
    const float* op = output_grad.as_fp32();
    const float* wp = weight.as_fp32();

#ifdef USE_CBLAS
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
        static_cast<int>(batch), static_cast<int>(in_dim), static_cast<int>(out_dim),
        1.0f, op, static_cast<int>(out_dim), wp, static_cast<int>(in_dim),
        0.0f, gp, static_cast<int>(in_dim));
#else
    for (size_t b = 0; b < batch; ++b) {
        for (size_t i = 0; i < in_dim; ++i) {
            float val = 0.0f;
            for (size_t o = 0; o < out_dim; ++o) {
                val += op[b * out_dim + o] * wp[o * in_dim + i];
            }
            gp[b * in_dim + i] = val;
        }
    }
#endif
    return grad;
}

static Tensor bias_backward(const Tensor& output_grad) {
    size_t out_dim = output_grad.shape_.back();
    size_t batch = output_grad.numel() / out_dim;

    Tensor grad({out_dim}, QuantType::FP32);
    float* gp = grad.as_fp32();
    const float* op = output_grad.as_fp32();

    memset(gp, 0, grad.data_size_);
    for (size_t b = 0; b < batch; ++b) {
        for (size_t o = 0; o < out_dim; ++o) {
            gp[o] += op[b * out_dim + o];
        }
    }
    return grad;
}

SwiGLUFFN::SwiGLUFFN(size_t d_model, size_t d_ff)
    : d_model_(d_model), d_ff_(d_ff > 0 ? d_ff : d_model * 4) {
    w_gate_ = std::make_shared<Linear>(d_model_, d_ff_, true);
    w_up_ = std::make_shared<Linear>(d_model_, d_ff_, true);
    w_down_ = std::make_shared<Linear>(d_ff_, d_model_, true);
}

Tensor SwiGLUFFN::forward(const Tensor& x) {
    cache_.input = x.clone();

    cache_.gate_out = w_gate_->forward(x);
    cache_.up_out = w_up_->forward(x);

    size_t n = cache_.gate_out.numel();
    cache_.gate_activated = cache_.gate_out.clone();

#ifdef USE_CUDA
    if (CudaContext::instance().is_available() && cache_.gate_activated.is_on_gpu()) {
        launch_silu(cache_.gate_activated.as_gpu_fp32(), n, CudaContext::instance().stream());
        cache_.gate_activated.gpu_dirty_ = true;
    } else {
#endif
    float* ga = cache_.gate_activated.as_fp32();
    for (size_t i = 0; i < n; ++i) {
        float v = ga[i];
        ga[i] = v / (1.0f + expf(-v));
    }
#ifdef USE_CUDA
    }
#endif

    cache_.multiplied = Tensor(cache_.gate_activated.shape_, QuantType::FP32);

#ifdef USE_CUDA
    if (CudaContext::instance().is_available() && cache_.gate_activated.is_on_gpu()) {
        cache_.multiplied.to_gpu();
        launch_elementwise_mul(cache_.multiplied.as_gpu_fp32(),
                    cache_.gate_activated.as_gpu_fp32(),
                    cache_.up_out.as_gpu_fp32(), n,
                    CudaContext::instance().stream());
        cache_.multiplied.gpu_dirty_ = true;
    } else {
#endif
    const float* ga_p = cache_.gate_activated.as_fp32();
    const float* up_p = cache_.up_out.as_fp32();
    float* mp = cache_.multiplied.as_fp32();
    for (size_t i = 0; i < n; ++i) {
        mp[i] = ga_p[i] * up_p[i];
    }
#ifdef USE_CUDA
    }
#endif

    return w_down_->forward(cache_.multiplied);
}

SwiGLUFFN::Gradients SwiGLUFFN::backward(const Tensor& output_grad) {
    Gradients grads;

    grads.w_down_weight_grad = linear_backward_weight(cache_.multiplied, output_grad);
    grads.w_down_bias_grad = bias_backward(output_grad);
    Tensor d_multiplied = linear_backward_input(output_grad, w_down_->weight);

    Tensor d_gate_activated(d_multiplied.shape_, QuantType::FP32);
    Tensor d_up(d_multiplied.shape_, QuantType::FP32);
    const float* dm = d_multiplied.as_fp32();
    const float* ga_p = cache_.gate_activated.as_fp32();
    const float* up_p = cache_.up_out.as_fp32();
    float* dga = d_gate_activated.as_fp32();
    float* dup = d_up.as_fp32();
    size_t n = d_multiplied.numel();
    for (size_t i = 0; i < n; ++i) {
        dga[i] = dm[i] * up_p[i];
        dup[i] = dm[i] * ga_p[i];
    }

    const float* go_p = cache_.gate_out.as_fp32();
    Tensor d_gate_out(d_gate_activated.shape_, QuantType::FP32);
    float* dgo = d_gate_out.as_fp32();
    for (size_t i = 0; i < n; ++i) {
        float sig = ga_p[i];
        dgo[i] = dga[i] * sig * (1.0f + go_p[i] * (1.0f - sig));
    }

    grads.w_gate_weight_grad = linear_backward_weight(cache_.input, d_gate_out);
    grads.w_gate_bias_grad = bias_backward(d_gate_out);
    Tensor d_input_gate = linear_backward_input(d_gate_out, w_gate_->weight);

    grads.w_up_weight_grad = linear_backward_weight(cache_.input, d_up);
    grads.w_up_bias_grad = bias_backward(d_up);
    Tensor d_input_up = linear_backward_input(d_up, w_up_->weight);

    grads.input_grad = Tensor(cache_.input.shape_, QuantType::FP32);
    const float* dig = d_input_gate.as_fp32();
    const float* diu = d_input_up.as_fp32();
    float* ig = grads.input_grad.as_fp32();
    for (size_t i = 0; i < grads.input_grad.numel(); ++i) {
        ig[i] = dig[i] + diu[i];
    }

    return grads;
}

#ifdef USE_CUDA
__global__ void kernel_silu_impl(float* data, size_t n) {
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float v = data[idx];
    data[idx] = v / (1.0f + expf(-v));
}

void launch_silu(float* data, size_t n, cudaStream_t stream) {
    int block = 256;
    int grid = (static_cast<int>(n) + block - 1) / block;
    kernel_silu_impl<<<grid, block, 0, stream>>>(data, n);
}

__global__ void kernel_elementwise_mul_impl(float* out, const float* a, const float* b, size_t n) {
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    out[idx] = a[idx] * b[idx];
}

void launch_elementwise_mul(float* out, const float* a, const float* b, size_t n, cudaStream_t stream) {
    int block = 256;
    int grid = (static_cast<int>(n) + block - 1) / block;
    kernel_elementwise_mul_impl<<<grid, block, 0, stream>>>(out, a, b, n);
}
#endif

} // namespace neuroflow