/*********************************************************************
  Blosc - Blocked Shuffling and Compression Library

  Copyright (c) 2021  Blosc Development Team <blosc@blosc.org>
  https://blosc.org
  License: BSD 3-Clause (see LICENSE.txt)

  See LICENSE.txt for details about copyright and rights to use.
**********************************************************************/

/*********************************************************************
  Bitshuffle - Filter for improving compression of typed binary data.

  Author: Kiyoshi Masui <kiyo@physics.ubc.ca>
  Website: https://github.com/kiyo-masui/bitshuffle

  Note: Adapted for c-blosc by Francesc Alted.

  See LICENSES/BITSHUFFLE.txt file for details about copyright and
  rights to use.
**********************************************************************/

#include "bitshuffle-avx2.h"
#include "bitshuffle-sse2.h"
#include "bitshuffle-generic.h"
#include <stdlib.h>

/* Make sure AVX2 is available for the compilation target and compiler. */
#if defined(__AVX2__)

#include <immintrin.h>

/* The next is useful for debugging purposes */
#if 0
#include <stdio.h>
#include <string.h>

static void printymm(__m256i ymm0)
{
  uint8_t buf[32];

  ((__m256i *)buf)[0] = ymm0;
  printf("%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x\n",
          buf[0], buf[1], buf[2], buf[3],
          buf[4], buf[5], buf[6], buf[7],
          buf[8], buf[9], buf[10], buf[11],
          buf[12], buf[13], buf[14], buf[15],
          buf[16], buf[17], buf[18], buf[19],
          buf[20], buf[21], buf[22], buf[23],
          buf[24], buf[25], buf[26], buf[27],
          buf[28], buf[29], buf[30], buf[31]);
}
#endif


/* ---- Code that requires AVX2. Intel Haswell (2013) and later. ---- */



/* Transpose bits within bytes. */
int64_t bshuf_trans_bit_byte_AVX(const void* in, void* out, const size_t size,
                                 const size_t elem_size) {

  size_t ii, kk;
  const char* in_b = (const char*) in;
  char* out_b = (char*) out;
  int32_t* out_i32;

  size_t nbyte = elem_size * size;

  int64_t count;

  __m256i ymm;
  int32_t bt;

  for (ii = 0; ii + 31 < nbyte; ii += 32) {
    ymm = _mm256_loadu_si256((__m256i *) &in_b[ii]);
    for (kk = 0; kk < 8; kk++) {
      bt = _mm256_movemask_epi8(ymm);
      ymm = _mm256_slli_epi16(ymm, 1);
      out_i32 = (int32_t*) &out_b[((7 - kk) * nbyte + ii) / 8];
      *out_i32 = bt;
    }
  }
  count = bshuf_trans_bit_byte_remainder(in, out, size, elem_size,
                                         nbyte - nbyte % 32);
  return count;
}


/* Transpose bits within elements. */
int64_t bshuf_trans_bit_elem_AVX(const void* in, void* out, const size_t size,
                                 const size_t elem_size) {

  int64_t count;

  CHECK_MULT_EIGHT(size);

  void* tmp_buf = malloc(size * elem_size);
  if (tmp_buf == NULL) return -1;

  count = bshuf_trans_byte_elem_SSE(in, out, size, elem_size);
  CHECK_ERR_FREE(count, tmp_buf);
  count = bshuf_trans_bit_byte_AVX(out, tmp_buf, size, elem_size);
  CHECK_ERR_FREE(count, tmp_buf);
  count = bshuf_trans_bitrow_eight(tmp_buf, out, size, elem_size);

  free(tmp_buf);

  return count;
}


/* For data organized into a row for each bit (8 * elem_size rows), transpose
 * the bytes. */
int64_t bshuf_trans_byte_bitrow_AVX(const void* in, void* out, const size_t size,
                                    const size_t elem_size) {

  size_t hh, ii, jj, kk, mm;
  const char* in_b = (const char*) in;
  char* out_b = (char*) out;

  CHECK_MULT_EIGHT(size);

  size_t nrows = 8 * elem_size;
  size_t nbyte_row = size / 8;

  if (elem_size % 4) return bshuf_trans_byte_bitrow_SSE(in, out, size,
                                                        elem_size);

  __m256i ymm_0[8];
  __m256i ymm_1[8];
  __m256i ymm_storeage[8][4];

  for (jj = 0; jj + 31 < nbyte_row; jj += 32) {
    for (ii = 0; ii + 3 < elem_size; ii += 4) {
      for (hh = 0; hh < 4; hh ++) {

        for (kk = 0; kk < 8; kk ++){
          ymm_0[kk] = _mm256_loadu_si256((__m256i *) &in_b[
              (ii * 8 + hh * 8 + kk) * nbyte_row + jj]);
        }

        for (kk = 0; kk < 4; kk ++){
          ymm_1[kk] = _mm256_unpacklo_epi8(ymm_0[kk * 2],
                                           ymm_0[kk * 2 + 1]);
          ymm_1[kk + 4] = _mm256_unpackhi_epi8(ymm_0[kk * 2],
                                               ymm_0[kk * 2 + 1]);
        }

        for (kk = 0; kk < 2; kk ++){
          for (mm = 0; mm < 2; mm ++){
            ymm_0[kk * 4 + mm] = _mm256_unpacklo_epi16(
                ymm_1[kk * 4 + mm * 2],
                ymm_1[kk * 4 + mm * 2 + 1]);
            ymm_0[kk * 4 + mm + 2] = _mm256_unpackhi_epi16(
                ymm_1[kk * 4 + mm * 2],
                ymm_1[kk * 4 + mm * 2 + 1]);
          }
        }

        for (kk = 0; kk < 4; kk ++){
          ymm_1[kk * 2] = _mm256_unpacklo_epi32(ymm_0[kk * 2],
                                                ymm_0[kk * 2 + 1]);
          ymm_1[kk * 2 + 1] = _mm256_unpackhi_epi32(ymm_0[kk * 2],
                                                    ymm_0[kk * 2 + 1]);
        }

        for (kk = 0; kk < 8; kk ++){
          ymm_storeage[kk][hh] = ymm_1[kk];
        }
      }

      for (mm = 0; mm < 8; mm ++) {

        for (kk = 0; kk < 4; kk ++){
          ymm_0[kk] = ymm_storeage[mm][kk];
        }

        ymm_1[0] = _mm256_unpacklo_epi64(ymm_0[0], ymm_0[1]);
        ymm_1[1] = _mm256_unpacklo_epi64(ymm_0[2], ymm_0[3]);
        ymm_1[2] = _mm256_unpackhi_epi64(ymm_0[0], ymm_0[1]);
        ymm_1[3] = _mm256_unpackhi_epi64(ymm_0[2], ymm_0[3]);

        ymm_0[0] = _mm256_permute2x128_si256(ymm_1[0], ymm_1[1], 32);
        ymm_0[1] = _mm256_permute2x128_si256(ymm_1[2], ymm_1[3], 32);
        ymm_0[2] = _mm256_permute2x128_si256(ymm_1[0], ymm_1[1], 49);
        ymm_0[3] = _mm256_permute2x128_si256(ymm_1[2], ymm_1[3], 49);

        _mm256_storeu_si256((__m256i *) &out_b[
            (jj + mm * 2 + 0 * 16) * nrows + ii * 8], ymm_0[0]);
        _mm256_storeu_si256((__m256i *) &out_b[
            (jj + mm * 2 + 0 * 16 + 1) * nrows + ii * 8], ymm_0[1]);
        _mm256_storeu_si256((__m256i *) &out_b[
            (jj + mm * 2 + 1 * 16) * nrows + ii * 8], ymm_0[2]);
        _mm256_storeu_si256((__m256i *) &out_b[
            (jj + mm * 2 + 1 * 16 + 1) * nrows + ii * 8], ymm_0[3]);
      }
    }
  }
  for (ii = 0; ii < nrows; ii ++ ) {
    for (jj = nbyte_row - nbyte_row % 32; jj < nbyte_row; jj ++) {
      out_b[jj * nrows + ii] = in_b[ii * nbyte_row + jj];
    }
  }
  return size * elem_size;
}


/* Shuffle bits within the bytes of eight element blocks. */
int64_t bshuf_shuffle_bit_eightelem_AVX(const void* in, void* out, const size_t size,
                                        const size_t elem_size) {

  CHECK_MULT_EIGHT(size);

  // With a bit of care, this could be written such that such that it is
  // in_buf = out_buf safe.
  const char* in_b = (const char*) in;
  char* out_b = (char*) out;

  size_t ii, jj, kk;
  size_t nbyte = elem_size * size;

  __m256i ymm;
  int32_t bt;

  if (elem_size % 4) {
    return bshuf_shuffle_bit_eightelem_SSE(in, out, size, elem_size);
  } else {
    for (jj = 0; jj + 31 < 8 * elem_size; jj += 32) {
      for (ii = 0; ii + 8 * elem_size - 1 < nbyte;
           ii += 8 * elem_size) {
        ymm = _mm256_loadu_si256((__m256i *) &in_b[ii + jj]);
        for (kk = 0; kk < 8; kk++) {
          bt = _mm256_movemask_epi8(ymm);
          ymm = _mm256_slli_epi16(ymm, 1);
          size_t ind = (ii + jj / 8 + (7 - kk) * elem_size);
          * (int32_t *) &out_b[ind] = bt;
        }
      }
    }
  }
  return size * elem_size;
}


/* Untranspose bits within elements. */
int64_t bshuf_untrans_bit_elem_AVX(const void* in, void* out, const size_t size,
                                   const size_t elem_size) {

  int64_t count;

  CHECK_MULT_EIGHT(size);

  void* tmp_buf = malloc(size * elem_size);
  if (tmp_buf == NULL) return -1;

  count = bshuf_trans_byte_bitrow_AVX(in, tmp_buf, size, elem_size);
  CHECK_ERR_FREE(count, tmp_buf);
  count =  bshuf_shuffle_bit_eightelem_AVX(tmp_buf, out, size, elem_size);

  free(tmp_buf);
  return count;
}

const bool is_bshuf_AVX = true;

#else /* defined(__AVX2__) */

const bool is_bshuf_AVX = false;

int64_t
bshuf_trans_bit_elem_AVX(const void* in, void* out, const size_t size,
                         const size_t elem_size) {
  abort();
}

int64_t
bshuf_untrans_bit_elem_AVX(const void* in, void* out, const size_t size,
                           const size_t elem_size) {
  abort();
}

#endif /* defined(__AVX2__) */
