/*
 * SPDX-License-Identifier: LicenseRef-CSSL-1.0
 */

// Initial unit test for nr_layer_precoder_simd_[cm, simd] functions
// Issue: gNB: SIMD Precoding Bug in nr_layer_precoder_simd for PMI = 4
// https://gitlab.eurecom.fr/oai/openairinterface5g/-/issues/955

#include "gtest/gtest.h"
#include <cstdint>
#include <simde/x86/avx512.h>

#ifdef __cplusplus
extern "C" {
#endif

#include "openair1/PHY/TOOLS/tools_defs.h" // c16_t
#include "nfapi/open-nFAPI/nfapi/public_inc/fapi_nr_ue_interface.h" // fapi_nr_dl_config_dci_dl_pdu_rel15_t
#include "common/utils/LOG/log.h"

//#include "openair1/PHY/MODULATION/nr_modulation.h" // nr_layer_precoder_simd_[cm, simd]

// Forward declare the function but hide the Variable-Length Array (VLA) feature from C not available in C++
c16_t nr_layer_precoder_cm(int n_layers,
                           int symbol_size,
                           const c16_t (*dataF_in)[/*symbol_size*/], // c16_t dataF_in[n_layers][symbol_size]
                           int ant,
                           c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS],
                           int offset);

// Forward declare the function but hide the Variable-Length Array (VLA) feature from C not available in C++
void nr_layer_precoder_simd(int n_layers,
                            int symbol_size,
                            const c16_t (*dataF_in)[/*symbol_size*/], // c16_t dataF_in[n_layers][symbol_size]
                            int ant,
                            c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS],
                            int offset,
                            int re_cnt,
                            c16_t *dataF_out);

// Forward declare the layer mapper while hiding its C VLA dimension.
void nr_layer_mapping(int nb_codewords,
                      const c16_t *const mod_symbols[2],
                      const uint32_t n_symbols_per_codeword[2],
                      uint8_t n_layers,
                      int layer_size,
                      c16_t (*tx_layers)[/*layer_size*/]);

#ifdef __cplusplus
}
#endif

TEST(NrLayerPrecoderTest, Basic)
{
  constexpr int n_layers = 2;
  constexpr int symbol_size = 24;
  constexpr int n_ants = 2;
  constexpr int ant = 1;
  constexpr int re_cnt = 24;

  // Initialize the 2D input data buffer
  std::vector<c16_t> buffer_in(n_layers * symbol_size);
  std::vector<c16_t> buffer_out(n_layers * symbol_size);

  for (int i = 0; i < n_layers * symbol_size; ++i) {
    buffer_in[i] = {static_cast<int16_t>(i + 1), static_cast<int16_t>(-i - 1)};
  }
  for (int i = 0; i < n_ants * symbol_size; ++i) {
    buffer_out[i] = {static_cast<int16_t>(0), static_cast<int16_t>(0)};
  }

  // Cast flat buffer to the required 2D Variable-Length Array (VLA) style pointer
  c16_t(*dataF_in)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_in.data());
  c16_t(*dataF_out)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_out.data());

  // Create and populate the weights
  c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  for (int layer = 0; layer < n_layers; ++layer) {
    weights[layer][ant] = (c16_t){-SHRT_MAX, 0};
  }

  // Call the C function
  for (int symbol = 0; symbol < re_cnt; symbol++)
    dataF_out[ant][symbol] = nr_layer_precoder_cm(n_layers, symbol_size, dataF_in, ant, weights, symbol);

  // You can calculate the expected value manually or test rough correctness
  for (int symbol = 0; symbol < re_cnt; symbol++) {
    int result_real = -26 - (symbol<<1);
    int result_imag = -24 - (symbol<<1);

    EXPECT_EQ(dataF_out[ant][symbol].r, result_real)
        << " at [" << ant << "][" << symbol << "] got real part: " << dataF_out[ant][symbol].r;
    EXPECT_EQ(dataF_out[ant][symbol].i, -result_imag)
        << " at [" << ant << "][" << symbol << "] got imag part: " << dataF_out[ant][symbol].i;
  }
}

TEST(NrLayerPrecoderTest, SIMD)
{
  constexpr int n_layers = 2;
  constexpr int symbol_size = 24;
  constexpr int n_ants = 2;
  constexpr int ant = 1;
  constexpr int re_cnt = 24;

  // Initialize the 2D input data buffer
  std::vector<c16_t> buffer_in(n_layers * symbol_size);
  std::vector<c16_t> buffer_out(n_layers * symbol_size);

  for (int i = 0; i < n_layers * symbol_size; ++i) {
    buffer_in[i] = {static_cast<int16_t>(i + 1), static_cast<int16_t>(-i - 1)};
  }
  for (int i = 0; i < n_ants * symbol_size; ++i) {
    buffer_out[i] = {static_cast<int16_t>(0), static_cast<int16_t>(0)};
  }

  // Cast flat buffer to the required 2D Variable-Length Array (VLA) style pointer
  c16_t(*dataF_in)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_in.data());
  c16_t(*dataF_out)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_out.data());

  // Create and populate the weights
  c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  for (int layer = 0; layer < n_layers; ++layer) {
    weights[layer][ant] = (c16_t){-SHRT_MAX, 0};
  }

  // Call the C function
  nr_layer_precoder_simd(n_layers, symbol_size, dataF_in, ant, weights, 0, re_cnt, dataF_out[ant]);

  // You can calculate the expected value manually or test rough correctness
  for (int symbol = 0; symbol < re_cnt; symbol++) {
    int result_real = -26 - (symbol<<1);
    int result_imag = -24 - (symbol<<1);

    EXPECT_EQ(dataF_out[ant][symbol].r, result_real)
        << " at [" << ant << "][" << symbol << "] got real part: " << dataF_out[ant][symbol].r;
    EXPECT_EQ(dataF_out[ant][symbol].i, -result_imag)
        << " at [" << ant << "][" << symbol << "] got imag part: " << dataF_out[ant][symbol].i;
  }
}

TEST(NrLayerPrecoderTest, Compare_CM_SIMD)
{
  constexpr int n_layers = 1;
  constexpr int symbol_size = 24;
  constexpr int n_ants = 2;
  constexpr int re_cnt = 24;

  // Initialize the 2D input data buffer
  std::vector<c16_t> buffer_in(n_layers * symbol_size);
  std::vector<c16_t> buffer_out_cm(n_ants * symbol_size);
  std::vector<c16_t> buffer_out_simd(n_ants * symbol_size);

  for (int i = 0; i < n_layers * symbol_size; ++i) {
	buffer_in[i] = {static_cast<int16_t>((rand() % (2 * SHRT_MAX + 1)) - SHRT_MAX), static_cast<int16_t>((rand() % (2 * SHRT_MAX + 1)) - SHRT_MAX)};
  }
  for (int i = 0; i < n_ants * symbol_size; ++i) {
    buffer_out_cm[i] = {static_cast<int16_t>(0), static_cast<int16_t>(0)};
    buffer_out_simd[i] = {static_cast<int16_t>(0), static_cast<int16_t>(0)};
  }

  // Cast flat buffer to the required 2D Variable-Length Array (VLA) style pointer
  c16_t(*dataF_in)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_in.data());
  c16_t(*dataF_out_cm)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_out_cm.data());
  c16_t(*dataF_out_simd)[symbol_size] = reinterpret_cast<c16_t(*)[symbol_size]>(buffer_out_simd.data());

  // Create and populate the weights
  c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  for (int layer = 0; layer < n_layers; ++layer) {
    // Could not use convert_precoder_weight() as complex.h could not be used in googletest
    // Use the logic in convert_precoder_weight()
    // precoder [−1,−j]
    weights[layer][0] = (c16_t){-SHRT_MAX, 0};
    weights[layer][1] = (c16_t){0, -SHRT_MAX};
  }

  // Get the results for all the antenna 
  for (int ant = 0; ant < n_ants; ant++) {

    // Call the C function
    for (int symbol = 0; symbol < re_cnt; symbol++)
      dataF_out_cm[ant][symbol] = nr_layer_precoder_cm(n_layers, symbol_size, dataF_in, ant, weights, symbol);

    // Call the C function
    nr_layer_precoder_simd(n_layers, symbol_size, dataF_in, ant, weights, 0, re_cnt, dataF_out_simd[ant]);

    // Compare the result from both C function
    for (int symbol = 0; symbol < re_cnt; symbol++) {
      EXPECT_EQ(dataF_out_cm[ant][symbol].r, dataF_out_simd[ant][symbol].r)
          << " at [" << ant << "][" << symbol << "] got real part: " << dataF_out_cm[ant][symbol].r << " result " << dataF_out_simd[ant][symbol].r;
      EXPECT_EQ(dataF_out_cm[ant][symbol].i, dataF_out_simd[ant][symbol].i)
          << " at [" << ant << "][" << symbol << "] got imag part: " << dataF_out_cm[ant][symbol].i << " result " << dataF_out_simd[ant][symbol].i;
    }
  }
}

TEST(NrLayerMappingTest, Rank8UsesTwoCodewordsAcrossEightDistinctLayers)
{
  constexpr int n_codewords = 2;
  constexpr int n_layers = 8;
  constexpr int layers_per_codeword = 4;
  constexpr int symbols_per_layer = 7;
  constexpr int symbols_per_codeword = layers_per_codeword * symbols_per_layer;

  std::vector<c16_t> cw0(symbols_per_codeword);
  std::vector<c16_t> cw1(symbols_per_codeword);
  for (int i = 0; i < symbols_per_codeword; ++i) {
    cw0[i] = {static_cast<int16_t>(100 + i), static_cast<int16_t>(-100 - i)};
    cw1[i] = {static_cast<int16_t>(1000 + i), static_cast<int16_t>(-1000 - i)};
  }

  const c16_t *mod_symbols[2] = {cw0.data(), cw1.data()};
  const uint32_t n_symbols_per_codeword[2] = {symbols_per_codeword, symbols_per_codeword};
  std::vector<c16_t> layer_storage(n_layers * symbols_per_layer, c16_t{SHRT_MIN, SHRT_MIN});
  c16_t(*tx_layers)[symbols_per_layer] = reinterpret_cast<c16_t(*)[symbols_per_layer]>(layer_storage.data());

  nr_layer_mapping(n_codewords,
                   mod_symbols,
                   n_symbols_per_codeword,
                   n_layers,
                   symbols_per_layer,
                   tx_layers);

  for (int re = 0; re < symbols_per_layer; ++re) {
    for (int layer = 0; layer < n_layers; ++layer) {
      const int cw = layer / layers_per_codeword;
      const int cw_layer = layer % layers_per_codeword;
      const c16_t expected = mod_symbols[cw][re * layers_per_codeword + cw_layer];
      EXPECT_EQ(tx_layers[layer][re].r, expected.r) << "real mismatch at layer " << layer << ", RE " << re;
      EXPECT_EQ(tx_layers[layer][re].i, expected.i) << "imag mismatch at layer " << layer << ", RE " << re;
    }
  }
}

int main(int argc, char **argv)
{
  // Initialize random seed
  srand(time(0));

  logInit();
  testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}