#include "../../../devices/bang/common_bang.h"
#include "../../../reduce/bang/reduce_bang.h"
#include "rms_norm_bang.h"

__nram__ char nram_buffer[NRAM_MAX_SIZE];

template <typename T, typename Tw>
__mlu_global__ void rmsnorm(T *output, const T *input, const Tw *weight,
                            size_t *shape, ptrdiff_t *output_strides, ptrdiff_t *input_strides,
                            float epsilon, int num_dims, int norm_dim_size) {
    // Calculate problem dimensions
    int batch_volume = 1;
    for (int dim_idx = 0; dim_idx < num_dims - 1; ++dim_idx) {
        batch_volume *= shape[dim_idx];
    }
    int vector_size = shape[num_dims - 1];

    // Task distribution across cores
    int remaining_tasks = batch_volume % taskDim;
    int base_tasks_per_core = batch_volume / taskDim;
    int actual_tasks = base_tasks_per_core + (taskId < remaining_tasks ? 1 : 0);
    int task_start_idx = (taskId < remaining_tasks ? taskId * (base_tasks_per_core + 1) : remaining_tasks * (base_tasks_per_core + 1) + (taskId - remaining_tasks) * base_tasks_per_core);

    // Determine optimal batch size based on vector size
    int max_batch_size;
    if (vector_size <= 64) {
        // For small vectors, process the entire vector at once
        max_batch_size = vector_size;
    } else {
        // For larger vectors, use optimized batch size
        max_batch_size = (NRAM_MAX_SIZE - 256) / (2 * sizeof(T) + sizeof(Tw) + sizeof(float));
        max_batch_size = std::min(max_batch_size, vector_size);
        max_batch_size = (max_batch_size / 64) * 64; // Align to 64 elements
    }

    constexpr int reduce_buffer_size = 128 / sizeof(float);

    // NRAM buffer allocation with dynamic sizing
    float *reduction_buffer = (float *)nram_buffer;
    T *input_cache = (T *)(reduction_buffer + reduce_buffer_size);
    Tw *weight_cache = (Tw *)(input_cache + max_batch_size);
    float *float_buffer = (float *)(weight_cache + max_batch_size);
    float *weight_float_buffer = (float *)(float_buffer + max_batch_size);

    // Process vectors assigned to current core
    for (int task_idx = 0; task_idx < actual_tasks; ++task_idx) {
        int current_index = task_start_idx + task_idx;

        // Calculate memory offsets for current task
        int input_offset = 0;
        int output_offset = 0;
        int temp_index = current_index;

        for (int dim = 0; dim < num_dims - 1; ++dim) {
            int dim_coord = temp_index % shape[dim];
            input_offset += dim_coord * input_strides[dim];
            output_offset += dim_coord * output_strides[dim];
            temp_index /= shape[dim];
        }

        // Compute sum of squares
        float sum_squared = 0.0f;

        if (vector_size <= 128) {
            // Small vector optimization: process entire vector at once
            __memcpy(input_cache, input + input_offset, vector_size * sizeof(T), GDRAM2NRAM);

            // Convert to float and square
            if constexpr (std::is_same<T, half>::value) {
                __bang_half2float(float_buffer, input_cache, vector_size);
            } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                __bang_bfloat162float(float_buffer, input_cache, vector_size);
            } else {
                __memcpy(float_buffer, input_cache, vector_size * sizeof(float), NRAM2NRAM);
            }

            __bang_mul(float_buffer, float_buffer, float_buffer, vector_size);

            // Direct accumulation for small vectors
            for (int i = 0; i < vector_size; ++i) {
                sum_squared += float_buffer[i];
            }
        } else {
            // Large vector processing with chunking
            __bang_write_zero(reduction_buffer, reduce_buffer_size);
            size_t processed_elements = 0;

            while (processed_elements < vector_size) {
                size_t current_batch = std::min((size_t)max_batch_size, vector_size - processed_elements);

                // Load input data
                __memcpy(input_cache, input + input_offset + processed_elements * input_strides[num_dims - 1],
                         current_batch * sizeof(T), GDRAM2NRAM);

                // Convert to float and square
                if constexpr (std::is_same<T, half>::value) {
                    __bang_half2float(float_buffer, input_cache, current_batch);
                } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                    __bang_bfloat162float(float_buffer, input_cache, current_batch);
                } else {
                    __memcpy(float_buffer, input_cache, current_batch * sizeof(float), NRAM2NRAM);
                }

                __bang_mul(float_buffer, float_buffer, float_buffer, current_batch);

                // Accumulate squared values
                float batch_sum = 0.0f;
                if (current_batch >= 128) {
                    op::common_bang::reduce_op::sumInternal(reduction_buffer, float_buffer, current_batch);
                    batch_sum = reduction_buffer[0];
                } else {
                    for (size_t i = 0; i < current_batch; ++i) {
                        batch_sum += float_buffer[i];
                    }
                }

                sum_squared += batch_sum;
                processed_elements += current_batch;
            }
        }

        // Compute normalization factor
        float rms_value = sqrtf(sum_squared / vector_size + epsilon);
        float inv_rms = 1.0f / rms_value;

        // Process vector for normalization
        if (vector_size <= max_batch_size) {
            // Process entire vector at once for small vectors
            __memcpy(input_cache, input + input_offset, vector_size * sizeof(T), GDRAM2NRAM);
            __memcpy(weight_cache, weight, vector_size * sizeof(Tw), GDRAM2NRAM);

            // Convert input to float
            if constexpr (std::is_same<T, half>::value) {
                __bang_half2float(float_buffer, input_cache, vector_size);
            } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                __bang_bfloat162float(float_buffer, input_cache, vector_size);
            } else {
                __memcpy(float_buffer, input_cache, vector_size * sizeof(float), NRAM2NRAM);
            }

            // Convert weight to float if needed
            if constexpr (std::is_same<Tw, half>::value) {
                __bang_half2float(weight_float_buffer, weight_cache, vector_size);
            } else if constexpr (std::is_same<Tw, bfloat16_t>::value) {
                __bang_bfloat162float(weight_float_buffer, weight_cache, vector_size);
            } else {
                __memcpy(weight_float_buffer, weight_cache, vector_size * sizeof(float), NRAM2NRAM);
            }

            // Multiply by weight and apply normalization
            __bang_mul(float_buffer, float_buffer, weight_float_buffer, vector_size);
            __bang_mul_scalar(float_buffer, float_buffer, inv_rms, vector_size);

            // Convert back to output type
            if constexpr (std::is_same<T, half>::value) {
                __bang_float2half(input_cache, float_buffer, vector_size);
            } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                __bang_float2bfloat16(input_cache, float_buffer, vector_size);
            } else {
                __memcpy(input_cache, float_buffer, vector_size * sizeof(float), NRAM2NRAM);
            }

            // Store results
            __memcpy(output + output_offset, input_cache, vector_size * sizeof(T), NRAM2GDRAM);
        } else {
            // Large vector processing with chunking
            size_t processed_elements = 0;
            while (processed_elements < vector_size) {
                size_t current_batch = std::min((size_t)max_batch_size, vector_size - processed_elements);

                // Load input and weight data
                __memcpy(input_cache, input + input_offset + processed_elements * input_strides[num_dims - 1],
                         current_batch * sizeof(T), GDRAM2NRAM);
                __memcpy(weight_cache, weight + processed_elements, current_batch * sizeof(Tw), GDRAM2NRAM);

                // Convert input to float
                if constexpr (std::is_same<T, half>::value) {
                    __bang_half2float(float_buffer, input_cache, current_batch);
                } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                    __bang_bfloat162float(float_buffer, input_cache, current_batch);
                } else {
                    __memcpy(float_buffer, input_cache, current_batch * sizeof(float), NRAM2NRAM);
                }

                // Convert weight to float if needed
                if constexpr (std::is_same<Tw, half>::value) {
                    __bang_half2float(weight_float_buffer, weight_cache, current_batch);
                } else if constexpr (std::is_same<Tw, bfloat16_t>::value) {
                    __bang_bfloat162float(weight_float_buffer, weight_cache, current_batch);
                } else {
                    __memcpy(weight_float_buffer, weight_cache, current_batch * sizeof(float), NRAM2NRAM);
                }

                // Multiply by weight and apply normalization
                __bang_mul(float_buffer, float_buffer, weight_float_buffer, current_batch);
                __bang_mul_scalar(float_buffer, float_buffer, inv_rms, current_batch);

                // Convert back to output type
                if constexpr (std::is_same<T, half>::value) {
                    __bang_float2half(input_cache, float_buffer, current_batch);
                } else if constexpr (std::is_same<T, bfloat16_t>::value) {
                    __bang_float2bfloat16(input_cache, float_buffer, current_batch);
                } else {
                    __memcpy(input_cache, float_buffer, current_batch * sizeof(float), NRAM2NRAM);
                }

                // Store results
                __memcpy(output + output_offset + processed_elements * output_strides[num_dims - 1],
                         input_cache, current_batch * sizeof(T), NRAM2GDRAM);

                processed_elements += current_batch;
            }
        }
    }
}

template <typename T, typename Tw>
void rmsnormUnion(void *workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, void *y, const void *x, const void *w, const size_t *shape, const ptrdiff_t *y_strides, const ptrdiff_t *x_strides, float eps, int ndim) {
    cnrtDim3_t kernel_dim;
    cnrtFunctionType_t kernel_type;

    // Configure kernel dimensions
    kernel_dim.x = core_per_cluster;
    kernel_dim.y = cluster_count;
    kernel_dim.z = 1;
    kernel_type = CNRT_FUNC_TYPE_UNION1; // Can choose others, but must adapt kernel_type accordingly
    int dimsize = shape[ndim - 1];       // Length of operation dimension
    int dim_s;                           // dim_s is the next power of 2 greater than dimsize
    float mi = log2(dimsize);
    if (floor(mi) == mi) {
        dim_s = dimsize;
    } else {
        dim_s = pow(2, floor(mi) + 1);
    }
    constexpr int reduce_num = 128 / sizeof(float); // Cambricon __bang_reduce_sum can only reduce 128 bytes at a time
    if (dim_s < reduce_num) {
        dim_s = reduce_num; // Force dim_s >= reduce_num
    }

    // Prepare device pointers
    auto y_ = reinterpret_cast<T *>(y);
    auto x_ = reinterpret_cast<const T *>(x);
    auto w_ = reinterpret_cast<const Tw *>(w);
    char *tmp_device = reinterpret_cast<char *>(workspace);
    char *tmp_stride = tmp_device + ndim * sizeof(size_t);
    size_t *mlu_shape = (size_t *)tmp_device;
    ptrdiff_t *mlu_x_strides = (ptrdiff_t *)tmp_stride;
    ptrdiff_t *mlu_y_strides = mlu_x_strides + ndim;

    // Copy shape and stride information to device
    CNRT_CHECK(cnrtMemcpyAsync(mlu_shape, const_cast<size_t *>(shape), ndim * sizeof(size_t), queue, cnrtMemcpyHostToDev)); // const not supported
    CNRT_CHECK(cnrtMemcpyAsync(mlu_x_strides, const_cast<ptrdiff_t *>(x_strides), ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev));
    CNRT_CHECK(cnrtMemcpyAsync(mlu_y_strides, const_cast<ptrdiff_t *>(y_strides), ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev));

    // Launch kernel
    rmsnorm<T, Tw><<<kernel_dim, kernel_type, queue>>>(y_, x_, w_, mlu_shape, mlu_y_strides, mlu_x_strides, eps, ndim, dim_s);

    cnrtQueueSync(queue);
}

namespace op::rms_norm::bang {

struct Descriptor::Opaque {
    std::shared_ptr<device::bang::Handle::Internal> internal;
};

Descriptor::~Descriptor() {
    delete _opaque;
}

infiniStatus_t Descriptor::create(
    infiniopHandle_t handle_,
    Descriptor **desc_ptr,
    infiniopTensorDescriptor_t y_desc,
    infiniopTensorDescriptor_t x_desc,
    infiniopTensorDescriptor_t w_desc,
    float epsilon) {
    auto handle = reinterpret_cast<device::bang::cambricon::Handle *>(handle_);

    auto result = RMSNormInfo::create(y_desc, x_desc, w_desc, epsilon);
    CHECK_RESULT(result);
    auto info = result.take();

    size_t workspace_size = info.ndim() * (sizeof(size_t) + 2 * sizeof(ptrdiff_t));
    *desc_ptr = new Descriptor(
        new Descriptor::Opaque{static_cast<device::bang::Handle *>(handle)->internal()},
        info,
        workspace_size,
        handle->device,
        handle->device_id);

    return INFINI_STATUS_SUCCESS;
}

infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size,
                                     void *y, const void *x, const void *w,
                                     void *stream) const {
    auto queue = reinterpret_cast<cnrtQueue_t>(stream);
    int core_per_cluster = _opaque->internal->getCorePerCluster();
    int cluster_count = _opaque->internal->getClusterCount();

    // Dispatch based on data types - support all combinations
    if (_info.atype == INFINI_DTYPE_F16) {
        if (_info.wtype == INFINI_DTYPE_F16) {
            rmsnormUnion<half, half>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_F32) {
            rmsnormUnion<half, float>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_BF16) {
            rmsnormUnion<half, bfloat16_t>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else {
            return INFINI_STATUS_BAD_TENSOR_DTYPE;
        }
    } else if (_info.atype == INFINI_DTYPE_F32) {
        if (_info.wtype == INFINI_DTYPE_F32) {
            rmsnormUnion<float, float>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_F16) {
            rmsnormUnion<float, half>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_BF16) {
            rmsnormUnion<float, bfloat16_t>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else {
            return INFINI_STATUS_BAD_TENSOR_DTYPE;
        }
    } else if (_info.atype == INFINI_DTYPE_BF16) {
        if (_info.wtype == INFINI_DTYPE_BF16) {
            rmsnormUnion<bfloat16_t, bfloat16_t>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_F32) {
            rmsnormUnion<bfloat16_t, float>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else if (_info.wtype == INFINI_DTYPE_F16) {
            rmsnormUnion<bfloat16_t, half>(workspace, core_per_cluster, cluster_count, queue, y, x, w, _info.shape.data(), _info.y_strides.data(), _info.x_strides.data(), _info.epsilon, _info.ndim());
        } else {
            return INFINI_STATUS_BAD_TENSOR_DTYPE;
        }
    } else {
        return INFINI_STATUS_BAD_TENSOR_DTYPE;
    }

    return INFINI_STATUS_SUCCESS;
}

} // namespace op::rms_norm::bang
