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

#include <climits>
#include <gtest/gtest.h>
extern "C" {
#include "nr_common.h"
#include "common/utils/LOG/log.h"
}

TEST(nr_common, nr_timer) {
  NR_timer_t timer;
  nr_timer_setup(&timer, 10, 1);
  nr_timer_start(&timer);
  EXPECT_TRUE(nr_timer_is_active(&timer));
  EXPECT_FALSE(nr_timer_expired(&timer));
  for (auto i = 0; i < 10; i++) {
    nr_timer_tick(&timer);
  }
  EXPECT_FALSE(nr_timer_is_active(&timer));
  EXPECT_TRUE(nr_timer_expired(&timer));
}

TEST(nr_common, pdsch_codeword_layer_split)
{
  constexpr uint8_t expected_codewords[] = {1, 1, 1, 1, 2, 2, 2, 2};
  constexpr uint8_t expected_cw0_layers[] = {1, 2, 3, 4, 2, 3, 3, 4};
  constexpr uint8_t expected_cw1_layers[] = {0, 0, 0, 0, 3, 3, 4, 4};

  for (uint8_t layers = 1; layers <= NR_MAX_NB_LAYERS; layers++) {
    const int index = layers - 1;
    EXPECT_EQ(nr_pdsch_num_codewords(layers), expected_codewords[index]);
    EXPECT_EQ(nr_pdsch_layers_per_codeword(layers, 0), expected_cw0_layers[index]);
    if (expected_codewords[index] == 2) {
      EXPECT_EQ(nr_pdsch_layers_per_codeword(layers, 1), expected_cw1_layers[index]);
    }
  }
}

TEST(nr_common, pdsch_two_codeword_ack_tracks_only_pending_transport_blocks)
{
  EXPECT_EQ(nr_pdsch_pending_after_ack(0, 3, 0), 3); // both NACK or DTX
  EXPECT_EQ(nr_pdsch_pending_after_ack(0, 3, 1), 2); // CW0 ACK, CW1 NACK
  EXPECT_EQ(nr_pdsch_pending_after_ack(0, 3, 2), 1); // CW0 NACK, CW1 ACK
  EXPECT_EQ(nr_pdsch_pending_after_ack(0, 3, 3), 0); // both ACK
  EXPECT_EQ(nr_pdsch_pending_after_ack(2, 3, 1), 2); // padding CW0 ACK is irrelevant
  EXPECT_EQ(nr_pdsch_pending_after_ack(2, 3, 2), 0); // pending CW1 ACK finishes
  EXPECT_EQ(nr_pdsch_pending_after_ack(1, 3, 2), 1); // padding CW1 ACK is irrelevant
  EXPECT_EQ(nr_pdsch_pending_after_ack(1, 3, 1), 0); // pending CW0 ACK finishes
  EXPECT_EQ(nr_pdsch_pending_after_ack(0, 1, 1), 0); // one-codeword legacy HARQ
}

TEST(nr_common, pdsch_rank_capacity_prefers_independent_modes)
{
  float eigenvalues[NR_MAX_NB_LAYERS];
  for (float &eigenvalue : eigenvalues)
    eigenvalue = 100.0f;
  float scores[NR_MAX_NB_LAYERS] = {};
  nr_pdsch_rank_capacity_scores(eigenvalues, NR_MAX_NB_LAYERS, 1.0f, scores);

  for (int rank = 1; rank < NR_MAX_NB_LAYERS; rank++)
    EXPECT_GT(scores[rank], scores[rank - 1]);
  EXPECT_EQ(nr_pdsch_select_rank(scores, NR_MAX_NB_LAYERS, 0.05f), 8);
}

TEST(nr_common, pdsch_rank_capacity_rejects_weak_modes)
{
  float eigenvalues[NR_MAX_NB_LAYERS] = {100.0f, 0.01f, 0.01f, 0.01f, 0.01f, 0.01f, 0.01f, 0.01f};
  float scores[NR_MAX_NB_LAYERS] = {};
  nr_pdsch_rank_capacity_scores(eigenvalues, NR_MAX_NB_LAYERS, 1.0f, scores);

  EXPECT_EQ(nr_pdsch_select_rank(scores, NR_MAX_NB_LAYERS, 0.05f), 1);
}

TEST(nr_common, pdsch_rank_capacity_respects_rank_and_gain_margin)
{
  const float scores[NR_MAX_NB_LAYERS] = {10.0f, 10.4f, 10.6f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f};
  EXPECT_EQ(nr_pdsch_select_rank(scores, 3, 0.05f), 3);
  EXPECT_EQ(nr_pdsch_select_rank(scores, 2, 0.05f), 1);
}

TEST(nr_common, type1_sp_4x1_parameter_ranges)
{
  constexpr uint8_t expected_i11[] = {16, 16, 16, 16, 16, 16, 8, 8};
  constexpr uint8_t expected_i13[] = {1, 4, 3, 3, 1, 1, 1, 1};
  constexpr uint8_t expected_i2[] = {4, 2, 2, 2, 2, 2, 2, 2};
  for (uint8_t layers = 1; layers <= NR_MAX_NB_LAYERS; layers++) {
    EXPECT_EQ(nr_type1_sp_4x1_i11_count(layers), expected_i11[layers - 1]);
    EXPECT_EQ(nr_type1_sp_4x1_i13_count(layers), expected_i13[layers - 1]);
    EXPECT_EQ(nr_type1_sp_4x1_i2_count(layers), expected_i2[layers - 1]);
  }

  c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  EXPECT_FALSE(nr_type1_sp_4x1_weights(0, 0, 0, 0, weights));
  EXPECT_FALSE(nr_type1_sp_4x1_weights(8, 8, 0, 0, weights));
  EXPECT_FALSE(nr_type1_sp_4x1_weights(8, 0, 1, 0, weights));
  EXPECT_FALSE(nr_type1_sp_4x1_weights(8, 0, 0, 2, weights));
}

TEST(nr_common, type1_sp_4x1_precoders_are_orthogonal)
{
  for (uint8_t layers = 1; layers <= NR_MAX_NB_LAYERS; layers++) {
    const uint8_t i11_count = nr_type1_sp_4x1_i11_count(layers);
    for (uint8_t i11 = 0; i11 < i11_count; i11++) {
      for (uint8_t i13 = 0; i13 < nr_type1_sp_4x1_i13_count(layers); i13++) {
        for (uint8_t i2 = 0; i2 < nr_type1_sp_4x1_i2_count(layers); i2++) {
          c16_t weights[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
          ASSERT_TRUE(nr_type1_sp_4x1_weights(layers, i11, i13, i2, weights));
          const int64_t expected_norm = 8LL * SHRT_MAX * SHRT_MAX / layers;
          const int64_t tolerance = expected_norm / 1000;
          for (int a = 0; a < layers; a++) {
            int64_t diagonal = 0;
            for (int port = 0; port < 8; port++)
              diagonal += (int64_t)weights[a][port].r * weights[a][port].r
                          + (int64_t)weights[a][port].i * weights[a][port].i;
            EXPECT_NEAR((double)diagonal, (double)expected_norm, (double)tolerance);
            for (int b = a + 1; b < layers; b++) {
              int64_t inner_re = 0;
              int64_t inner_im = 0;
              for (int port = 0; port < 8; port++) {
                inner_re += (int64_t)weights[a][port].r * weights[b][port].r
                            + (int64_t)weights[a][port].i * weights[b][port].i;
                inner_im += (int64_t)weights[a][port].r * weights[b][port].i
                            - (int64_t)weights[a][port].i * weights[b][port].r;
              }
              EXPECT_NEAR((double)inner_re, 0.0, (double)tolerance);
              EXPECT_NEAR((double)inner_im, 0.0, (double)tolerance);
            }
          }
        }
      }
    }
  }
}

TEST(nr_common, type1_sp_4x1_i2_rotates_cross_polarization)
{
  c16_t i2_zero[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  c16_t i2_one[NR_MAX_NB_LAYERS][NR_MAX_CSI_PORTS];
  ASSERT_TRUE(nr_type1_sp_4x1_weights(8, 0, 0, 0, i2_zero));
  ASSERT_TRUE(nr_type1_sp_4x1_weights(8, 0, 0, 1, i2_one));
  EXPECT_EQ(i2_zero[0][4].i, 0);
  EXPECT_EQ(i2_one[0][4].r, 0);
  EXPECT_EQ(i2_zero[0][4].r, i2_one[0][4].i);
}


int main(int argc, char **argv)
{
  logInit();
  testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}