/********************************************************************************
 * Copyright (c) 2025 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include "aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp"

#include <aidge/operator/Memorize.hpp>
#include <aidge/utils/Types.h>
#include <memory>
#include <cstdlib>
#include <random>

#include <catch2/catch_test_macros.hpp>

#include "aidge/data/Tensor.hpp"
#include "aidge/backend/cpu/operator/HeavisideImpl.hpp"
#include "aidge/graph/Node.hpp"
#include "aidge/utils/TensorUtils.hpp"

#include "aidge/operator/Add.hpp"

namespace Aidge
{

TEST_CASE("[cpu/operator] Heaviside(forward)", "[Heaviside][CPU]") {

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> valueDist(-1.0f, 1.0f);
    std::uniform_int_distribution<std::size_t> dimSizeDist(std::size_t(2), std::size_t(10));
    std::uniform_int_distribution<std::size_t> nbDimsDist(std::size_t(1), std::size_t(5));

    SECTION("1D Tensor") {

        std::shared_ptr<Tensor> input0 = std::make_shared<Tensor>(Array1D<float,10> {
            {0, 1, 2,-3, 4,-5,-6, 7, 8, 9}
        });
        std::shared_ptr<Tensor> expectedOutput = std::make_shared<Tensor>(Array1D<float,10> {
            {0.5, 1, 1, 0, 1, 0, 0, 1, 1, 1}
        });

        std::shared_ptr<Node> heaviside = Heaviside(0.5);
        auto op = std::static_pointer_cast<OperatorTensor>(heaviside->getOperator());
        op->associateInput(0, input0);
        op->setBackend("cpu");
        op->setDataType(DataType::Float32);

        op->forward();
        REQUIRE(approxEq<float>(*op->getOutput(0),*expectedOutput));
    }

    SECTION("+1-D Tensor")
    {
        auto dims = std::vector<std::size_t>();
        auto nbDims = nbDimsDist(gen);

        for (auto i = 0u; i < nbDims; ++i) {
            dims.push_back(dimSizeDist(gen));
        }

        auto numberOfElements = std::accumulate(dims.cbegin(), dims.cend(), std::size_t(1), std::multiplies<std::size_t>());
        float* inputArray = new float[numberOfElements];
        float* resultArray = new float[numberOfElements];

        for(auto i = 0u; i < numberOfElements; ++i)
        {
            inputArray[i] = valueDist(gen);
            resultArray[i] = inputArray[i] > 0 ? 1 : (inputArray[i] == 0 ? 0.5 : 0);
        }

        auto T0 = std::make_shared<Tensor>();
        T0->setDataType(DataType::Float32);
        T0->setBackend("cpu");

        auto T1 = std::make_shared<Tensor>();
        T1->setDataType(DataType::Float32);
        T1->setBackend("cpu");

        T0->resize(dims);
        T0->getImpl()->setRawPtr(inputArray, numberOfElements);
        T1->resize(dims);
        T1->getImpl()->setRawPtr(resultArray, numberOfElements);

        std::shared_ptr<Node> heaviside = Heaviside(0.5);
        auto op = std::static_pointer_cast<OperatorTensor>(heaviside->getOperator());
        op->associateInput(0, T0);
        op->setBackend("cpu");
        op->setDataType(DataType::Float32);

        op->forward();

        REQUIRE(approxEq<float>(*(op->getOutput(0)), *T1));
    }
}

TEST_CASE("[cpu/operator] Heaviside(backward)", "[Heaviside][CPU]") {
    auto hs = Heaviside(1.0f);
    auto op = std::static_pointer_cast<OperatorTensor>(hs->getOperator());
    op->setDataType(DataType::Float32);
    op->setBackend("cpu");

    auto input = Tensor(Array1D<float, 3>({1.0, -1.0, 1.0}));
    input.setDataType(DataType::Float32);
    input.setBackend("cpu");

    auto grad = Tensor(Array1D<float, 3>({1.0, 1.0, 1.0}));
    grad.setDataType(DataType::Float32);
    grad.setBackend("cpu");

    op->setInput(IOIndex_t(0), std::make_shared<Tensor>(input));
    op->forward();

    Log::info("Output : ");
    op->getOutput(0)->print();

    op->getOutput(0)->setGrad(std::make_shared<Tensor>(grad));
    op->backward();

    Log::info("Gradient : ");
    op->getInput(0)->grad()->print();
}
}