/*
 * Classic McEliece Matrix Logic
 * Based on the public domain reference implementation by the designers
 * (https://classic.mceliece.org/impl.html - released in Oct 2022 for NISTPQC-R4)
 *
 *
 * (C) 2023 Jack Lloyd
 *     2023,2024 Fabian Albert, Amos Treiber - Rohde & Schwarz Cybersecurity
 *
 * Botan is released under the Simplified BSD License (see license.txt)
 **/

#include <botan/internal/cmce_matrix.h>

#include <botan/strong_type.h>
#include <botan/internal/buffer_slicer.h>
#include <botan/internal/buffer_stuffer.h>

namespace Botan {

namespace {

// Strong types for matrix used internally by Classic_McEliece_Matrix
using CmceMatrixRow = Strong<secure_bitvector, struct CmceMatrixRow_>;
using CmceMatrix = Strong<std::vector<CmceMatrixRow>, struct CmceMatrix_>;

}  // Anonymous namespace

namespace {

CT::Mask<uint64_t> bit_at_mask(uint64_t val, size_t pos) {
   return CT::Mask<uint64_t>::expand((static_cast<uint64_t>(1) << pos) & val);
}

/// Swaps bit i with bit j in val
void swap_1_bit(uint64_t& val, size_t i, size_t j) {
   const uint64_t bit_i = (val >> i) & CT::value_barrier<uint64_t>(1);
   const uint64_t bit_j = (val >> j) & CT::value_barrier<uint64_t>(1);
   const uint64_t xor_sum = bit_i ^ bit_j;
   val ^= (xor_sum << i);
   val ^= (xor_sum << j);
}

size_t count_lsb_zeros(uint64_t n) {
   size_t res = 0;
   auto found_only_zeros = Botan::CT::Mask<uint64_t>::set();
   for(size_t bit_pos = 0; bit_pos < sizeof(uint64_t) * 8; ++bit_pos) {
      auto bit_set_mask = bit_at_mask(n, bit_pos);
      found_only_zeros &= ~bit_set_mask;
      res += static_cast<size_t>(found_only_zeros.if_set_return(1));
   }

   return res;
}

CmceMatrix init_matrix_with_alphas(const Classic_McEliece_Parameters& params,
                                   const Classic_McEliece_Field_Ordering& field_ordering,
                                   const Classic_McEliece_Minimal_Polynomial& g) {
   auto alphas = field_ordering.alphas(params.n());
   std::vector<Classic_McEliece_GF> inv_g_of_alpha;
   inv_g_of_alpha.reserve(params.n());
   for(const auto& alpha : alphas) {
      inv_g_of_alpha.push_back(g(alpha).inv());
   }
   CmceMatrix mat(std::vector<CmceMatrixRow>(params.pk_no_rows(), CmceMatrixRow(params.n())));

   for(size_t i = 0; i < params.t(); ++i) {
      for(size_t j = 0; j < params.n(); ++j) {
         const auto inv_g = inv_g_of_alpha[j].elem().get();
         for(size_t alpha_i_j_bit = 0; alpha_i_j_bit < params.m(); ++alpha_i_j_bit) {
            mat[i * params.m() + alpha_i_j_bit][j] =
               CT::Choice::from_int(static_cast<uint16_t>((inv_g >> alpha_i_j_bit) & 1));
         }
      }
      // Update for the next i so that:
      // inv_g_of_alpha[j] = h_i_j = alpha_j^i/g(alpha_j)
      for(size_t j = 0; j < params.n(); ++j) {
         inv_g_of_alpha.at(j) *= alphas.at(j);
      }
   }

   return mat;
}

std::optional<CmceColumnSelection> move_columns(CmceMatrix& mat, const Classic_McEliece_Parameters& params) {
   BOTAN_ASSERT(mat.size() == params.pk_no_rows(), "Matrix has incorrect number of rows");
   BOTAN_ASSERT(mat.get().at(0).size() == params.n(), "Matrix has incorrect number of columns");
   static_assert(Classic_McEliece_Parameters::nu() == 64, "nu needs to be 64");

   const size_t pos_offset = params.pk_no_rows() - Classic_McEliece_Parameters::mu();

   // Get the area of the matrix that needs to be (potentially) swapped.
   // Its the sub m*t x nu matrix at column m*t - mu. For const time reasons,
   // the sub-matrix is represented as an array of uint64_ts, where the 1st
   // bit is the least significant bit
   std::vector<uint64_t> matrix_swap_area;
   matrix_swap_area.reserve(params.pk_no_rows());
   for(size_t i = 0; i < params.pk_no_rows(); ++i) {
      matrix_swap_area.push_back(mat[i].subvector<uint64_t>(pos_offset));
   }

   // To find which columns need to be swapped to allow for a systematic matrix form, we need to
   // investigate how a gauss algorithm affects the last mu rows of the swap area.
   std::array<uint64_t, Classic_McEliece_Parameters::mu()> sub_mat;  // NOLINT(*-member-init)

   // Extract the bottom mu x nu matrix at offset pos_offset
   for(size_t i = 0; i < Classic_McEliece_Parameters::mu(); i++) {
      sub_mat[i] = matrix_swap_area[pos_offset + i];
   }

   std::array<size_t, Classic_McEliece_Parameters::mu()> pivot_indices = {0};  // ctz_list

   // Identify the pivot indices, i.e., the indices of the leading ones for all rows
   // when transforming the matrix into semi-systematic form. This algorithm is a modified
   // Gauss algorithm.
   for(size_t row_idx = 0; row_idx < Classic_McEliece_Parameters::mu(); ++row_idx) {
      // Identify pivots (index of first 1) by OR-ing all subsequent rows into row_acc
      auto row_acc = sub_mat.at(row_idx);
      for(size_t next_row = row_idx + 1; next_row < Classic_McEliece_Parameters::mu(); ++next_row) {
         row_acc |= sub_mat.at(next_row);
      }

      auto semi_systematic_form_failed = CT::Mask<uint64_t>::is_zero(row_acc);
      if(semi_systematic_form_failed.as_choice().as_bool()) {
         // If the current row and all subsequent rows are zero
         // we cannot create a semi-systematic matrix
         return std::nullopt;
      }

      // Using the row accumulator we can predict the index of the pivot
      // bit for the current row, i.e., the first index where we can set
      // the bit to one row by adding any subsequent row
      const size_t current_pivot_idx = count_lsb_zeros(row_acc);
      pivot_indices.at(row_idx) = current_pivot_idx;

      // Add subsequent rows to the current row, until the pivot
      // bit is set.
      for(size_t next_row = row_idx + 1; next_row < Classic_McEliece_Parameters::mu(); ++next_row) {
         // Add next row if the pivot bit is still zero
         auto add_next_row_mask = ~bit_at_mask(sub_mat.at(row_idx), current_pivot_idx);
         sub_mat.at(row_idx) ^= add_next_row_mask.if_set_return(sub_mat.at(next_row));
      }

      // Add the (new) current row to all subsequent rows, where the leading
      // bit of the current bit is one. Therefore, the column of the leading
      // bit becomes zero.
      // Note: In normal gauss, we would also add the current row to rows
      //       above the current one. However, here we only need to identify
      //       the columns to swap. Therefore, we can ignore the upper rows.
      for(size_t next_row = row_idx + 1; next_row < Classic_McEliece_Parameters::mu(); ++next_row) {
         // Add the current row to next_row if the pivot bit of next_row is set
         auto add_to_next_row_mask = bit_at_mask(sub_mat.at(next_row), current_pivot_idx);
         sub_mat.at(next_row) ^= add_to_next_row_mask.if_set_return(sub_mat.at(row_idx));
      }
   }

   // Create pivot bitvector from the pivot index vector
   CmceColumnSelection pivots(Classic_McEliece_Parameters::nu());
   for(auto pivot_idx : pivot_indices) {
      for(size_t i = 0; i < Classic_McEliece_Parameters::nu(); ++i) {
         auto mask_is_at_current_idx = Botan::CT::Mask<size_t>::is_equal(i, pivot_idx);
         pivots.at(i) = mask_is_at_current_idx.as_choice() || pivots.at(i).as_choice();
      }
   }

   // Swap the rows so the matrix can be transformed into systematic form
   for(size_t mat_row = 0; mat_row < params.pk_no_rows(); ++mat_row) {
      for(size_t col = 0; col < Classic_McEliece_Parameters::mu(); ++col) {
         swap_1_bit(matrix_swap_area.at(mat_row), col, pivot_indices.at(col));
      }
   }

   // Reinsert the swapped columns into the matrix
   for(size_t row = 0; row < params.pk_no_rows(); ++row) {
      mat[row].subvector_replace(pos_offset, matrix_swap_area[row]);
   }

   return pivots;
}

std::optional<CmceColumnSelection> apply_gauss(const Classic_McEliece_Parameters& params, CmceMatrix& mat) {
   BOTAN_ASSERT(mat.size() == params.pk_no_rows(), "Matrix has incorrect number of rows");
   BOTAN_ASSERT(mat.get().at(0).size() == params.n(), "Matrix has incorrect number of columns");
   // Initialized for systematic form instances
   // Is overridden for semi systematic instances
   auto pivots = CmceColumnSelection({0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0});

   // Gaussian Elimination
   for(size_t diag_pos = 0; diag_pos < params.pk_no_rows(); ++diag_pos) {
      if(params.is_f() && diag_pos == params.pk_no_rows() - params.mu()) {
         auto ret_pivots = move_columns(mat, params);
         const bool move_columns_failed = !ret_pivots.has_value();
         CT::unpoison(move_columns_failed);
         if(move_columns_failed) {
            return std::nullopt;
         } else {
            pivots = std::move(ret_pivots.value());
         }
      }

      // Iterates over all rows next_row under row diag_pos. If the bit at column
      // diag_pos differs between row diag_pos and row next_row, row next_row is added to row diag_pos.
      // This achieves that the respective bit at the diagonal becomes 1
      // (if mat is systematic)
      for(size_t next_row = diag_pos + 1; next_row < params.pk_no_rows(); ++next_row) {
         mat[diag_pos].get().ct_conditional_xor(!mat[diag_pos].at(diag_pos).as_choice(), mat[next_row].get());
      }

      // If the current bit on the diagonal is not set at this point
      // the matrix is not systematic. We abort the computation in this case.
      const bool diag_bit_zero = !mat[diag_pos].at(diag_pos);
      CT::unpoison(diag_bit_zero);
      if(diag_bit_zero) {
         return std::nullopt;
      }

      // Now the new row is added to all other rows, where the
      // bit in the column of the current position on the diagonal
      // is still one
      for(size_t row = 0; row < params.pk_no_rows(); ++row) {
         if(row != diag_pos) {
            mat[row].get().ct_conditional_xor(mat[row].at(diag_pos).as_choice(), mat[diag_pos].get());
         }
      }
   }

   return pivots;
}

std::vector<uint8_t> extract_pk_bytes_from_matrix(const Classic_McEliece_Parameters& params, const CmceMatrix& mat) {
   // Store T of the matrix (I_mt|T) as a linear vector to represent the
   // public key as defined in McEliece ISO 9.2.7
   std::vector<uint8_t> big_t(params.pk_size_bytes());
   auto big_t_stuffer = BufferStuffer(big_t);

   for(size_t row = 0; row < params.pk_no_rows(); ++row) {
      mat[row].subvector(params.pk_no_rows()).to_bytes(big_t_stuffer.next(params.pk_row_size_bytes()));
   }

   BOTAN_ASSERT_NOMSG(big_t_stuffer.full());

   return big_t;
}

}  // namespace

std::optional<std::pair<Classic_McEliece_Matrix, CmceColumnSelection>> Classic_McEliece_Matrix::create_matrix(
   const Classic_McEliece_Parameters& params,
   const Classic_McEliece_Field_Ordering& field_ordering,
   const Classic_McEliece_Minimal_Polynomial& g) {
   auto mat = init_matrix_with_alphas(params, field_ordering, g);
   auto pivots = apply_gauss(params, mat);

   auto gauss_failed = !pivots.has_value();
   CT::unpoison(gauss_failed);
   if(gauss_failed) {
      return std::nullopt;
   }

   auto pk_mat_bytes = extract_pk_bytes_from_matrix(params, mat);
   return std::make_pair(Classic_McEliece_Matrix(params, std::move(pk_mat_bytes)), pivots.value());
}

std::optional<std::pair<Classic_McEliece_Matrix, CmceColumnSelection>>
Classic_McEliece_Matrix::create_matrix_and_apply_pivots(const Classic_McEliece_Parameters& params,
                                                        Classic_McEliece_Field_Ordering& field_ordering,
                                                        const Classic_McEliece_Minimal_Polynomial& g) {
   auto pk_matrix_and_pivots = create_matrix(params, field_ordering, g);

   const bool matrix_creation_failed = !pk_matrix_and_pivots.has_value();
   CT::unpoison(matrix_creation_failed);
   if(matrix_creation_failed) {
      return std::nullopt;
   }

   auto& [_, pivots] = pk_matrix_and_pivots.value();

   if(params.is_f()) {
      field_ordering.permute_with_pivots(params, pivots);
   }

   return pk_matrix_and_pivots;
}

CmceCodeWord Classic_McEliece_Matrix::mul(const Classic_McEliece_Parameters& params, const CmceErrorVector& e) const {
   auto s = e.subvector<CmceCodeWord>(0, params.pk_no_rows());
   auto e_T = e.subvector(params.pk_no_rows());
   auto pk_slicer = BufferSlicer(m_mat_bytes);

   for(size_t i = 0; i < params.pk_no_rows(); ++i) {
      auto pk_current_bytes = pk_slicer.take(params.pk_row_size_bytes());
      auto row = secure_bitvector(pk_current_bytes, params.n() - params.pk_no_rows());
      row &= e_T;
      s[i] ^= row.has_odd_hamming_weight();
   }

   BOTAN_ASSERT_NOMSG(pk_slicer.empty());
   return s;
}
}  // namespace Botan
