/*
    SWARM

    Copyright (C) 2012-2025 Torbjorn Rognes and Frederic Mahe

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as
    published by the Free Software Foundation, either version 3 of the
    License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

    Contact: Torbjorn Rognes <torognes@ifi.uio.no>,
    Department of Informatics, University of Oslo,
    PO Box 1080 Blindern, NO-0316 Oslo, Norway
*/

#include "utils/pseudo_rng.h"
#include "zobrist.h"
#include <algorithm> // std::for_each
#include <cassert>
#include <cstddef>  // std::ptrdiff_t
#include <cstdint>  // uint64_t
#include <iterator>  // std::next
#include <vector>


// anonymous namespace: limit visibility and usage to this translation unit
namespace {

  auto to_uchar(char const nucleotide) -> unsigned char {
    // note: compressed nucleotides are in the range [-127, +127]
    return static_cast<unsigned char>(nucleotide);
  }

}  // end of anonymous namespace


uint64_t * zobrist_tab_base = nullptr;
uint64_t * zobrist_tab_byte_base = nullptr;


auto fill_rng_table(const unsigned int zobrist_len,
                    std::vector<uint64_t> & zobrist_tab_base_v) -> void
{
  /*
    Generate 4n random 64-bit numbers. They will represent the four
    different bases in any position (1 to n) of a sequence.  They will
    be XOR'ed together to form the hash of that sequence.  The number
    n, here 'zobrist_len', should be the length of the longest
    sequence to be hashed including potential additional insertions.

    The number is generated by xor'ing together four shifted
    31-bit random numbers.
  */
  static constexpr auto multiplier = 16U;

  /* allocate base table and fill with random 64 bit numbers */
  zobrist_tab_base_v.resize(4ULL * zobrist_len);
  zobrist_tab_base = zobrist_tab_base_v.data();

  std::for_each(zobrist_tab_base_v.begin(),
                zobrist_tab_base_v.end(),
                [](uint64_t & rng_value) -> void {
                  // refactoring: comment states 31-bit random numbers?!
                  rng_value = rand_64();
                  rng_value <<= multiplier;
                  rng_value ^= rand_64();
                  rng_value <<= multiplier;
                  rng_value ^= rand_64();
                  rng_value <<= multiplier;
                  rng_value ^= rand_64();
                });
}


auto fill_rng_byte_table(const unsigned int zobrist_len,
                         std::vector<uint64_t> const & zobrist_tab_base_v,
                         std::vector<uint64_t> & zobrist_tab_byte_base_v) -> void
{
  static constexpr auto byte_range = 256U;

  /* allocate byte table and combine into bytes for faster computations */
  zobrist_tab_byte_base_v.resize(1ULL * byte_range * (zobrist_len / 4));
  zobrist_tab_byte_base = zobrist_tab_byte_base_v.data();

  for(auto i = 0U; i < zobrist_len / 4; ++i) {
    for(auto j = 0U; j < byte_range; ++j) {
      auto rng_value = 0ULL;
      auto offset = j;
      // rng value stored at: 4 *  position   +  offset & 3U (= 0, 1, 2, or 3)
      rng_value ^= zobrist_tab_base_v[(4 * (4 * i + 0)) + (offset & 3U)];
      offset >>= 2U;
      rng_value ^= zobrist_tab_base_v[(4 * (4 * i + 1)) + (offset & 3U)];
      offset >>= 2U;
      rng_value ^= zobrist_tab_base_v[(4 * (4 * i + 2)) + (offset & 3U)];
      offset >>= 2U;
      rng_value ^= zobrist_tab_base_v[(4 * (4 * i + 3)) + (offset & 3U)];
      zobrist_tab_byte_base_v[(byte_range * i) + j] = rng_value;
    }
  }
}


auto zobrist_init(const unsigned int zobrist_len,
                  std::vector<uint64_t> & zobrist_tab_base_v,
                  std::vector<uint64_t> & zobrist_tab_byte_base_v) -> void
{
  fill_rng_table(zobrist_len, zobrist_tab_base_v);
  fill_rng_byte_table(zobrist_len, zobrist_tab_base_v, zobrist_tab_byte_base_v);
}


auto zobrist_exit() -> void
{
  zobrist_tab_byte_base = nullptr;
  zobrist_tab_base = nullptr;
}


auto zobrist_value(const unsigned int pos, const unsigned char offset) -> uint64_t
{
  assert(offset == 0 or offset == 1 or offset == 2 or offset == 3);
  return *std::next(zobrist_tab_base, (4 * pos) + offset);
}


auto zobrist_hash(char const * seq, unsigned int const len) -> uint64_t
{
  /* compute the Zobrist hash function of sequence seq of length len. */
  /* len is the actual number of bases in the sequence */
  /* it is encoded in (len + 3 ) / 4 bytes */

  // refactoring: equivalent to a std::reduce() algorithm?
  static constexpr auto offset = 64U;
  static constexpr auto nt_per_uint64 = 32U;  // 32 nucleotides can fit in a uint64
  uint64_t zobrist_hash = 0;
  auto pos = 0U;

  while (pos + nt_per_uint64 < len)  // hash bytes by groups of 8
    {
      for(auto i = 0U; i < nt_per_uint64; i += 4) {
        auto const a_byte = to_uchar(*seq);
        auto const target_hash = (offset * (pos + i)) + a_byte;
        assert(target_hash <= std::numeric_limits<std::ptrdiff_t>::max());
        auto const target_hash_signed = static_cast<std::ptrdiff_t>(target_hash);
        // i = {0, 4, 8, 12, 16, 20, 24, 28}
        zobrist_hash ^= *std::next(zobrist_tab_byte_base, target_hash_signed);
        seq = std::next(seq);
      }
      pos += nt_per_uint64;
    }

  while (pos + 4 < len)  // less than 8 bytes remaining (hash bytes one-by-one)
    {
      auto const a_byte = to_uchar(*seq);
      auto const target_hash = (offset * pos) + a_byte;
      assert(target_hash <= std::numeric_limits<std::ptrdiff_t>::max());
      auto const target_hash_signed = static_cast<std::ptrdiff_t>(target_hash);
      zobrist_hash ^= *std::next(zobrist_tab_byte_base, target_hash_signed);
      seq = std::next(seq);
      pos += 4;
    }

  if (pos < len)  // less than a byte remaining (hash pairs of bits one-by-one)
    {
      auto last_byte = to_uchar(*seq);
      seq = std::next(seq);  // useless in normal conditions, keep it for now
      while (pos < len)
        {
          zobrist_hash ^= zobrist_value(pos, last_byte & 3U);
          last_byte >>= 2U;
          ++pos;
        }
    }

  return zobrist_hash;
}


// refactoring: unrolling does not work (hard to deal with the last positions (sub-byte)
// refactoring: factorize zobrist_hash_delete_first() and zobrist_hash_insert_first()?
auto zobrist_hash_delete_first(char const * seq, unsigned int const len) -> uint64_t
{
  /* compute the Zobrist hash function of sequence seq,
     but delete the first base */

  static constexpr auto nt_per_byte = 4U;  // 4 nucleotides per byte
  static constexpr auto divider = 2U;
  auto offset = to_uchar(*seq);
  uint64_t zobrist_hash = 0;
  for(auto pos = 1U; pos < len; ++pos)
    {
      auto const is_new_byte = (pos & (nt_per_byte - 1)) == 0;
      if (is_new_byte) {  // every 4 positions, except the first one
        auto const target_chunk = (pos >> divider);
        assert(target_chunk <= std::numeric_limits<std::ptrdiff_t>::max());
        auto const target_chunk_signed = static_cast<std::ptrdiff_t>(target_chunk);
        offset = to_uchar(*std::next(seq, target_chunk_signed));  // load new data every 4 positions
      }
      else {
        offset >>= 2U;
      }
      zobrist_hash ^= zobrist_value(pos - 1, offset & 3U);
    }
  return zobrist_hash;
}


auto zobrist_hash_insert_first(char const * seq, const unsigned int len) -> uint64_t
{
  /* compute the Zobrist hash function of sequence seq,
     but insert a gap (no value) before the first base */

  static constexpr auto nt_per_byte = 4U;  // 4 nucleotides per byte
  static constexpr auto divider = 2U;
  auto offset = to_uchar(*seq);
  uint64_t zobrist_hash = 0;
  for(auto pos = 0U; pos < len; ++pos)
    {
      auto const is_new_byte = (pos & (nt_per_byte - 1)) == 0;
      if (is_new_byte) {  // every 4 positions, except the first one
        auto const target_chunk = (pos >> divider);
        assert(target_chunk <= std::numeric_limits<std::ptrdiff_t>::max());
        auto const target_chunk_signed = static_cast<std::ptrdiff_t>(target_chunk);
        offset = to_uchar(*std::next(seq, target_chunk_signed));  // load new data every 4 positions
      }
      else {
        offset >>= 2U;
      }
      zobrist_hash ^= zobrist_value(pos + 1, offset & 3U);
    }
  return zobrist_hash;
}
