/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#include "squared_difference_aicpu.h"

#include "cpu_kernel_utils.h"
#include "utils/eigen_tensor.h"
#include "utils/kernel_util.h"

namespace {
const uint32_t kOutputNum = 1;
const uint32_t kInputNum = 2;
const char* const kSquaredDifference = "SquaredDifference";
// when input data size is more than kParallelDataNum, use Parallel func
const int64_t kParallelDataNum = 2 * 1024;
const int64_t kParallelDataNumMid = 32 * 1024;
const int64_t kParallelDataNumSameShape = 7 * 1024;
const int64_t kParallelDataNumSameShapeMid = 35 * 1024;

#define SQUAREDDIFFERENCE_COMPUTE_CASE(DTYPE, TYPE, CTX)                  \
    case (DTYPE): {                                                       \
        uint32_t result = SquaredDifferenceCompute<TYPE>(CTX);            \
        if (result != KERNEL_STATUS_OK) {                                 \
            KERNEL_LOG_ERROR("SquaredDifference kernel compute failed."); \
            return result;                                                \
        }                                                                 \
        break;                                                            \
    }
} // namespace

namespace aicpu {
uint32_t SquaredDifferenceCpuKernel::Compute(CpuKernelContext& ctx)
{
    KERNEL_HANDLE_ERROR(
        NormalCheck(ctx, kInputNum, kOutputNum), "SquaredDifference check input and output number failed.");
    KERNEL_HANDLE_ERROR(SquaredDifferenceCheck(ctx), "SquaredDifference check params failed.");
    DataType data_type = ctx.Input(0)->GetDataType();
    switch (data_type) {
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_FLOAT16, Eigen::half, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_FLOAT, float, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_DOUBLE, double, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_INT32, int32_t, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_INT64, int64_t, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_COMPLEX64, std::complex<float>, ctx)
        SQUAREDDIFFERENCE_COMPUTE_CASE(DT_COMPLEX128, std::complex<double>, ctx)
        default:
            KERNEL_LOG_ERROR("SquaredDifference kernel data type [%s] not support.", DTypeStr(data_type).c_str());
            return KERNEL_STATUS_PARAM_INVALID;
    }
    return KERNEL_STATUS_OK;
}

template <typename T>
auto CalcDiffByType(T diff) -> T
{
    return diff;
}

template <>
std::complex<float> CalcDiffByType<std::complex<float>>(std::complex<float> diff)
{
    return std::conj(diff);
}

template <>
std::complex<double> CalcDiffByType<std::complex<double>>(std::complex<double> diff)
{
    return std::conj(diff);
}

uint32_t SquaredDifferenceCpuKernel::SquaredDifferenceCheck(const CpuKernelContext& ctx) const
{
    // the non null of input_0, input_1, output has been verified in NormalCheck
    Tensor* input_0 = ctx.Input(0);
    Tensor* input_1 = ctx.Input(1);
    Tensor* output = ctx.Output(0);
    KERNEL_CHECK_NULLPTR(input_0->GetData(), KERNEL_STATUS_PARAM_INVALID, "Get input 0 data failed.")
    KERNEL_CHECK_NULLPTR(input_1->GetData(), KERNEL_STATUS_PARAM_INVALID, "Get input 1 data failed.")
    KERNEL_CHECK_NULLPTR(output->GetData(), KERNEL_STATUS_PARAM_INVALID, "Get output data failed")
    DataType input0_type = input_0->GetDataType();
    DataType input1_type = input_1->GetDataType();
    KERNEL_CHECK_FALSE(
        (input0_type == input1_type), KERNEL_STATUS_PARAM_INVALID,
        "The data type of input0 [%s] need be same with "
        "input1 [%s].",
        DTypeStr(input0_type).c_str(), DTypeStr(input1_type).c_str())
    KERNEL_LOG_DEBUG(
        "SquaredDifferenceCpuKernel[%s], input0: size[%lu];"
        "input1: size[%lu], output: size[%lu].",
        ctx.GetOpType().c_str(), input_0->GetDataSize(), input_1->GetDataSize(), output->GetDataSize());

    return KERNEL_STATUS_OK;
}

// special compute is used in the following situations.
// 1. the shapes of input1 and input2 are the same
// 2. input1 is a 1D tensor with only one element or input1 is scalar
// 3. input2 is a 1D tensor with only one element or input2 is scalar
// 4. the shapes of input1 and input2 are different
template <typename T>
void SquaredDifferenceCpuKernel::SpecialCompute(
    BcastShapeType type, int64_t start, int64_t end, const T* input1, const T* input2, T* output)
{
    switch (type) {
        case BcastShapeType::SAME_SHAPE:
            for (int64_t i = start; i < end; ++i) {
                auto diff = *(input1 + i) - *(input2 + i);
                *(output + i) = diff * CalcDiffByType(diff);
            }
            break;
        case BcastShapeType::X_ONE_ELEMENT:
            for (int64_t i = start; i < end; ++i) {
                T temp = *input1;
                auto diff = (temp) - *(input2 + i);
                *(output + i) = diff * CalcDiffByType(diff);
            }
            break;
        case BcastShapeType::Y_ONE_ELEMENT:
            for (int64_t i = start; i < end; ++i) {
                T temp = *input2;
                auto diff = *(input1 + i) - (temp);
                *(output + i) = diff * CalcDiffByType(diff);
            }
            break;
        default:
            KERNEL_LOG_WARN("Invalid type [%d]", static_cast<int32_t>(type));
            break;
    }
}

template <typename T>
uint32_t SquaredDifferenceCpuKernel::NoBcastCompute(const CpuKernelContext& ctx)
{
    auto in0 = static_cast<T*>(ctx.Input(0)->GetData());
    auto in1 = static_cast<T*>(ctx.Input(1)->GetData());
    auto out = static_cast<T*>(ctx.Output(0)->GetData());
    int64_t in0_elements_nums = ctx.Input(0)->NumElements();
    int64_t in1_elements_nums = ctx.Input(1)->NumElements();
    int64_t data_num = ctx.Output(0)->NumElements();
    BcastShapeType type = in0_elements_nums == in1_elements_nums ?
                              BcastShapeType::SAME_SHAPE :
                              (in0_elements_nums == 1 ? BcastShapeType::X_ONE_ELEMENT : BcastShapeType::Y_ONE_ELEMENT);

    if (data_num >= kParallelDataNumSameShape) {
        uint32_t min_core_num = 1;
        int64_t max_core_num = std::max(min_core_num, aicpu::CpuKernelUtils::GetCPUNum(ctx) - kResvCpuNum);

        if (data_num <= kParallelDataNumSameShapeMid) {
            max_core_num = std::min(max_core_num, static_cast<int64_t>(4)); // up to 4 cpu cores
        }

        if (max_core_num > data_num) {
            max_core_num = data_num;
        }

        auto sharder_squareddifference = [this, &type, &in0, &in1, &out](int64_t start, int64_t end) {
            SpecialCompute<T>(type, start, end, in0, in1, out);
        };

        if (static_cast<int>(max_core_num) == 0) {
            return KERNEL_STATUS_PARAM_INVALID;
        }
        KERNEL_HANDLE_ERROR(
            CpuKernelUtils::ParallelFor(ctx, data_num, data_num / max_core_num, sharder_squareddifference),
            "SquaredDifference Compute failed.")
    } else {
        SpecialCompute<T>(type, 0, data_num, in0, in1, out);
    }

    return KERNEL_STATUS_OK;
}

template <typename T>
uint32_t SquaredDifferenceCpuKernel::BcastCompute(const CpuKernelContext& ctx, const Bcast& bcast)
{
    auto in0 = static_cast<T*>(ctx.Input(0)->GetData());
    auto in1 = static_cast<T*>(ctx.Input(1)->GetData());
    auto out = static_cast<T*>(ctx.Output(0)->GetData());
    int64_t data_num = ctx.Output(0)->NumElements();
    if (data_num >= kParallelDataNum) {
        uint32_t min_core_num = 1;
        int64_t max_core_num = std::max(min_core_num, aicpu::CpuKernelUtils::GetCPUNum(ctx) - kResvCpuNum);

        if (data_num <= kParallelDataNumMid) {
            max_core_num = std::min(max_core_num, static_cast<int64_t>(4)); // up to 4 cpu cores
        }

        if (max_core_num > data_num) {
            max_core_num = data_num;
        }

        auto sharder_squareddifference = [this, &in0, &in1, &out, &bcast](int64_t start, int64_t end) {
            for (int64_t i = start; i < end; ++i) {
                auto diff = *(in0 + bcast.GetBroadcastXIndex(i)) - *(in1 + bcast.GetBroadcastYIndex(i));
                *(out + i) = diff * CalcDiffByType(diff);
            }
        };

        if (static_cast<int>(max_core_num) == 0) {
            return KERNEL_STATUS_PARAM_INVALID;
        }
        KERNEL_HANDLE_ERROR(
            CpuKernelUtils::ParallelFor(ctx, data_num, data_num / max_core_num, sharder_squareddifference),
            "SquaredDifference Compute failed.")
    } else {
        for (int64_t i = 0; i < data_num; ++i) {
            auto diff = *(in0 + bcast.GetBroadcastXIndex(i)) - *(in1 + bcast.GetBroadcastYIndex(i));
            *(out + i) = diff * CalcDiffByType(diff);
        }
    }
    return KERNEL_STATUS_OK;
}

template <typename T>
uint32_t SquaredDifferenceCpuKernel::SquaredDifferenceCompute(const CpuKernelContext& ctx)
{
    Tensor* input0_tensor = ctx.Input(0);
    auto input0_shape = input0_tensor->GetTensorShape()->GetDimSizes();
    int64_t input0_elements_nums = input0_tensor->NumElements();

    Tensor* input1_tensor = ctx.Input(1);
    auto input1_shape = input1_tensor->GetTensorShape()->GetDimSizes();
    int64_t input1_elements_nums = input1_tensor->NumElements();

    bool is_no_need_bcast =
        (input0_shape == input1_shape) || (input0_elements_nums == 1) || (input1_elements_nums == 1);
    if (is_no_need_bcast) {
        return NoBcastCompute<T>(ctx);
    } else {
        Bcast bcast(input0_shape, input1_shape);
        if (!bcast.IsValid()) {
            KERNEL_LOG_ERROR("[%s] broadcast failed.", ctx.GetOpType().c_str());
            return KERNEL_STATUS_PARAM_INVALID;
        }

        return BcastCompute<T>(ctx, bcast);
    }
}

REGISTER_CPU_KERNEL(kSquaredDifference, SquaredDifferenceCpuKernel);
} // namespace aicpu