// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT

#include <gtest/gtest.h>
#include <hip/hip_runtime.h>
#include <memory>
#include <vector>

#include <hipdnn_frontend.hpp>
#include <hipdnn_frontend/node/LayerNormNode.hpp>
#include <hipdnn_test_sdk/constants/LayernormConstants.hpp>
#include <hipdnn_test_sdk/utilities/IntegrationTestFixture.hpp>
#include <hipdnn_test_sdk/utilities/LiftingTestHelpers.hpp>
#include <hipdnn_test_sdk/utilities/TestUtilities.hpp>
#include <hipdnn_test_sdk/utilities/TestableGraph.hpp>
#include <hipdnn_test_sdk/utilities/ToVec.hpp>

using namespace hipdnn_frontend;
using namespace hipdnn_frontend::graph;
using hipdnn_tests::IntegrationTestFixture;
using hipdnn_tests::liftGraph;
using hipdnn_tests::liftGraphWithoutFinalization;
using hipdnn_tests::TestableGraphLifting;
using hipdnn_tests::toVec;
namespace ln_constants = hipdnn_tests::constants;

namespace
{
class IntegrationLayerNormDescriptorLifting : public IntegrationTestFixture
{
protected:
    // Builds a standard training layernorm graph with all tensors set
    static std::shared_ptr<TestableGraphLifting> buildTrainingGraph()
    {
        auto graph = std::make_shared<TestableGraphLifting>();
        graph->set_name("LayernormLiftingTest")
            .set_compute_data_type(DataType::FLOAT)
            .set_intermediate_data_type(DataType::FLOAT)
            .set_io_data_type(DataType::FLOAT);

        auto x = std::make_shared<TensorAttributes>();
        x->set_uid(ln_constants::K_LAYERNORM_TENSOR_X_UID)
            .set_name("X")
            .set_data_type(DataType::FLOAT)
            .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_X_DIMS))
            .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_X_STRIDES));

        auto scale = std::make_shared<TensorAttributes>();
        scale->set_uid(ln_constants::K_LAYERNORM_TENSOR_SCALE_UID)
            .set_name("SCALE")
            .set_data_type(DataType::FLOAT)
            .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_DIMS))
            .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_STRIDES));

        auto bias = std::make_shared<TensorAttributes>();
        bias->set_uid(ln_constants::K_LAYERNORM_TENSOR_BIAS_UID)
            .set_name("BIAS")
            .set_data_type(DataType::FLOAT)
            .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_DIMS))
            .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_STRIDES));

        auto epsilon = std::make_shared<TensorAttributes>();
        epsilon->set_uid(ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID)
            .set_name("EPSILON")
            .set_data_type(DataType::FLOAT)
            .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_EPSILON_DIMS))
            .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_EPSILON_STRIDES))
            .set_value(1e-5f);

        LayernormAttributes layernormAttrs;
        layernormAttrs.set_name("layernorm_op");
        layernormAttrs.set_forward_phase(NormFwdPhase::TRAINING);
        layernormAttrs.set_epsilon(epsilon);

        auto [y, mean, invVariance] = graph->layernorm(x, scale, bias, layernormAttrs);
        y->set_uid(ln_constants::K_LAYERNORM_TENSOR_Y_UID).set_output(true).set_name("Y");
        mean->set_uid(ln_constants::K_LAYERNORM_TENSOR_MEAN_UID).set_output(true).set_name("MEAN");
        invVariance->set_uid(ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID)
            .set_output(true)
            .set_name("INV_VARIANCE");

        return graph;
    }
};

// Builds a layernorm graph in TRAINING mode (with mean/inv_variance), lowers
// via build_operation_graph(handle), lifts back with fromBackendDescriptor(),
// and verifies all operation attributes and tensors.
TEST_F(IntegrationLayerNormDescriptorLifting, LayernormTrainingRoundTripViaCApi)
{
    auto graph = buildTrainingGraph();

    auto liftedGraph = liftGraph(*graph, _handle);
    ASSERT_NE(liftedGraph, nullptr);

    // Verify graph-level data types
    EXPECT_EQ(liftedGraph->get_compute_data_type(), DataType::FLOAT);
    EXPECT_EQ(liftedGraph->get_intermediate_data_type(), DataType::FLOAT);
    EXPECT_EQ(liftedGraph->get_io_data_type(), DataType::FLOAT);

    // Verify tensors by UID (7 tensors: x, scale, bias, epsilon, y, mean, inv_variance)
    auto tensorMap = liftedGraph->getTensorsByUid();
    ASSERT_EQ(tensorMap.size(), 7u);

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_X_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_name(), "X");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_X_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_X_STRIDES));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_data_type(), DataType::FLOAT);

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_SCALE_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_SCALE_UID]->get_name(), "SCALE");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_SCALE_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_SCALE_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_STRIDES));

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_BIAS_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_BIAS_UID]->get_name(), "BIAS");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_BIAS_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_BIAS_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_STRIDES));

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID]->get_name(), "EPSILON");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_EPSILON_DIMS));

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_Y_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_name(), "Y");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_Y_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_Y_STRIDES));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_data_type(), DataType::FLOAT);

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_MEAN_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_MEAN_UID]->get_name(), "MEAN");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_MEAN_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_MEAN_DIMS));

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID), 0u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID]->get_name(),
              "INV_VARIANCE");
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_DIMS));

    // Verify the lifted graph has 1 layernorm sub-node
    auto& subNodes = liftedGraph->getSubNodes();
    ASSERT_EQ(subNodes.size(), 1u);

    auto* lnNode = dynamic_cast<LayerNormNode*>(subNodes[0].get());
    ASSERT_NE(lnNode, nullptr);

    EXPECT_EQ(lnNode->attributes.get_forward_phase(), NormFwdPhase::TRAINING);
    EXPECT_EQ(lnNode->attributes.get_name(), "layernorm_op");
    // X_DIMS={2,64,32,32}, SCALE_DIMS={1,64,32,32}: 3 trailing dims are normalized
    EXPECT_EQ(lnNode->attributes.get_normalized_dim_count(), 3);

    // Verify tensor references on the node
    EXPECT_EQ(lnNode->attributes.get_x()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_X_UID);
    EXPECT_EQ(lnNode->attributes.get_scale()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_SCALE_UID);
    EXPECT_EQ(lnNode->attributes.get_bias()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_BIAS_UID);
    EXPECT_EQ(lnNode->attributes.get_epsilon()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID);
    EXPECT_EQ(lnNode->attributes.get_y()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_Y_UID);
    ASSERT_NE(lnNode->attributes.get_mean(), nullptr);
    EXPECT_EQ(lnNode->attributes.get_mean()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_MEAN_UID);
    ASSERT_NE(lnNode->attributes.get_inv_variance(), nullptr);
    EXPECT_EQ(lnNode->attributes.get_inv_variance()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID);
}

// Builds a layernorm graph in INFERENCE mode (no mean/inv_variance), lowers
// via build_operation_graph(handle), lifts back with fromBackendDescriptor(),
// and verifies all operation attributes and that optional tensors are absent.
TEST_F(IntegrationLayerNormDescriptorLifting, LayernormInferenceRoundTripViaCApi)
{
    auto graph = std::make_shared<TestableGraphLifting>();
    graph->set_name("LayernormInferenceLiftingTest")
        .set_compute_data_type(DataType::FLOAT)
        .set_intermediate_data_type(DataType::FLOAT)
        .set_io_data_type(DataType::FLOAT);

    auto x = std::make_shared<TensorAttributes>();
    x->set_uid(ln_constants::K_LAYERNORM_TENSOR_X_UID)
        .set_name("X")
        .set_data_type(DataType::FLOAT)
        .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_X_DIMS))
        .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_X_STRIDES));

    auto scale = std::make_shared<TensorAttributes>();
    scale->set_uid(ln_constants::K_LAYERNORM_TENSOR_SCALE_UID)
        .set_name("SCALE")
        .set_data_type(DataType::FLOAT)
        .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_DIMS))
        .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_STRIDES));

    auto bias = std::make_shared<TensorAttributes>();
    bias->set_uid(ln_constants::K_LAYERNORM_TENSOR_BIAS_UID)
        .set_name("BIAS")
        .set_data_type(DataType::FLOAT)
        .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_DIMS))
        .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_BIAS_STRIDES));

    auto epsilon = std::make_shared<TensorAttributes>();
    epsilon->set_uid(ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID)
        .set_name("EPSILON")
        .set_data_type(DataType::FLOAT)
        .set_dim(toVec(ln_constants::K_LAYERNORM_TENSOR_EPSILON_DIMS))
        .set_stride(toVec(ln_constants::K_LAYERNORM_TENSOR_EPSILON_STRIDES))
        .set_value(1e-5f);

    LayernormAttributes layernormAttrs;
    layernormAttrs.set_name("layernorm_inference_op");
    layernormAttrs.set_forward_phase(NormFwdPhase::INFERENCE);
    layernormAttrs.set_epsilon(epsilon);

    auto [y, mean, invVariance] = graph->layernorm(x, scale, bias, layernormAttrs);
    y->set_uid(ln_constants::K_LAYERNORM_TENSOR_Y_UID).set_output(true).set_name("Y");
    // mean and invVariance should be nullptr in INFERENCE mode
    EXPECT_EQ(mean, nullptr);
    EXPECT_EQ(invVariance, nullptr);

    auto liftedGraph = liftGraph(*graph, _handle);
    ASSERT_NE(liftedGraph, nullptr);

    // Verify tensors by UID (5 tensors: x, scale, bias, epsilon, y — no mean or inv_variance)
    auto tensorMap = liftedGraph->getTensorsByUid();
    ASSERT_EQ(tensorMap.size(), 5u);

    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_X_UID), 0u);
    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_SCALE_UID), 0u);
    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_BIAS_UID), 0u);
    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID), 0u);
    ASSERT_NE(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_Y_UID), 0u);

    // Mean and inv_variance should NOT be in the tensor map
    EXPECT_EQ(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_MEAN_UID), 0u);
    EXPECT_EQ(tensorMap.count(ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID), 0u);

    // Verify the lifted graph has 1 layernorm sub-node
    auto& subNodes = liftedGraph->getSubNodes();
    ASSERT_EQ(subNodes.size(), 1u);

    auto* lnNode = dynamic_cast<LayerNormNode*>(subNodes[0].get());
    ASSERT_NE(lnNode, nullptr);

    EXPECT_EQ(lnNode->attributes.get_forward_phase(), NormFwdPhase::INFERENCE);
    EXPECT_EQ(lnNode->attributes.get_name(), "layernorm_inference_op");

    // Optional tensors should be null on the lifted node
    EXPECT_EQ(lnNode->attributes.get_mean(), nullptr);
    EXPECT_EQ(lnNode->attributes.get_inv_variance(), nullptr);
}

// Verifies that tensor objects are shared between the tensor map and the
// layernorm node attributes after lifting.
TEST_F(IntegrationLayerNormDescriptorLifting, LayernormTensorSharingPreserved)
{
    auto graph = buildTrainingGraph();
    graph->set_name("LayernormTensorSharingTest");

    auto liftedGraph = liftGraph(*graph, _handle);
    ASSERT_NE(liftedGraph, nullptr);

    auto tensorMap = liftedGraph->getTensorsByUid();
    auto& subNodes = liftedGraph->getSubNodes();
    ASSERT_EQ(subNodes.size(), 1u);

    auto* lnNode = dynamic_cast<LayerNormNode*>(subNodes[0].get());
    ASSERT_NE(lnNode, nullptr);

    // Verify tensor objects are shared (same pointer) between map and node
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID].get(),
              lnNode->attributes.get_x().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_SCALE_UID].get(),
              lnNode->attributes.get_scale().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_BIAS_UID].get(),
              lnNode->attributes.get_bias().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID].get(),
              lnNode->attributes.get_epsilon().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID].get(),
              lnNode->attributes.get_y().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_MEAN_UID].get(),
              lnNode->attributes.get_mean().get());
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID].get(),
              lnNode->attributes.get_inv_variance().get());
}

// Builds a layernorm graph, serializes to binary, creates a backend descriptor
// from bytes (no handle, no finalize), calls fromBackendDescriptor(), and verifies
// the layernorm operation survives the FlatBuffer-direct deserialization path.
TEST_F(IntegrationLayerNormDescriptorLifting, LayernormLiftWithoutFinalization)
{
    auto graph = buildTrainingGraph();
    graph->set_name("LayernormFlatBufferLiftTest");

    auto liftedGraph = liftGraphWithoutFinalization(*graph);
    ASSERT_NE(liftedGraph, nullptr);

    EXPECT_EQ(liftedGraph->get_compute_data_type(), DataType::FLOAT);

    auto& subNodes = liftedGraph->getSubNodes();
    ASSERT_EQ(subNodes.size(), 1u);

    auto* lnNode = dynamic_cast<LayerNormNode*>(subNodes[0].get());
    ASSERT_NE(lnNode, nullptr);

    EXPECT_EQ(lnNode->attributes.get_forward_phase(), NormFwdPhase::TRAINING);
    EXPECT_EQ(lnNode->attributes.get_name(), "layernorm_op");
    EXPECT_EQ(lnNode->attributes.get_x()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_X_UID);
    EXPECT_EQ(lnNode->attributes.get_scale()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_SCALE_UID);
    EXPECT_EQ(lnNode->attributes.get_bias()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_BIAS_UID);
    EXPECT_EQ(lnNode->attributes.get_epsilon()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_EPSILON_UID);
    EXPECT_EQ(lnNode->attributes.get_y()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_Y_UID);
    ASSERT_NE(lnNode->attributes.get_mean(), nullptr);
    EXPECT_EQ(lnNode->attributes.get_mean()->get_uid(), ln_constants::K_LAYERNORM_TENSOR_MEAN_UID);
    ASSERT_NE(lnNode->attributes.get_inv_variance(), nullptr);
    EXPECT_EQ(lnNode->attributes.get_inv_variance()->get_uid(),
              ln_constants::K_LAYERNORM_TENSOR_INV_VARIANCE_UID);

    // Verify tensor dims
    auto tensorMap = liftedGraph->getTensorsByUid();
    ASSERT_EQ(tensorMap.size(), 7u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_X_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_X_STRIDES));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_Y_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_Y_UID]->get_stride(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_Y_STRIDES));
}

// Exercises the deserialize() path with a handle for a layernorm graph.
TEST_F(IntegrationLayerNormDescriptorLifting, LayernormDeserializeViaBackendWithHandle)
{
    auto graph = buildTrainingGraph();
    graph->set_name("LayernormDeserializeTest");

    auto result = graph->validate();
    ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg;

    auto [data, serErr] = graph->to_binary();
    ASSERT_TRUE(serErr.is_good()) << serErr.get_message();

    // Create a new graph and use deserialize with handle
    auto liftedGraph = std::make_shared<TestableGraphLifting>();
    result = liftedGraph->deserialize(_handle, data);
    ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg;

    // Verify graph-level data types
    EXPECT_EQ(liftedGraph->get_compute_data_type(), DataType::FLOAT);
    EXPECT_EQ(liftedGraph->get_intermediate_data_type(), DataType::FLOAT);
    EXPECT_EQ(liftedGraph->get_io_data_type(), DataType::FLOAT);

    // Verify the lifted graph has 1 operation node
    auto& subNodes = liftedGraph->getSubNodes();
    ASSERT_EQ(subNodes.size(), 1u);

    auto* lnNode = dynamic_cast<LayerNormNode*>(subNodes[0].get());
    ASSERT_NE(lnNode, nullptr);

    EXPECT_EQ(lnNode->attributes.get_forward_phase(), NormFwdPhase::TRAINING);
    EXPECT_EQ(lnNode->attributes.get_name(), "layernorm_op");

    // Verify tensor dims
    auto tensorMap = liftedGraph->getTensorsByUid();
    ASSERT_EQ(tensorMap.size(), 7u);
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_X_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_X_DIMS));
    EXPECT_EQ(tensorMap[ln_constants::K_LAYERNORM_TENSOR_SCALE_UID]->get_dim(),
              toVec(ln_constants::K_LAYERNORM_TENSOR_SCALE_DIMS));
}

} // namespace