/*
Copyright 2016 Esri

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

A local copy of the license and additional notices are located with the
source distribution at:

http://github.com/Esri/lepcc/

Contributors:  Thomas Maurer
*/

#include <algorithm>
#include "BitStuffer2.h"

using namespace std;
using namespace lepcc;

// -------------------------------------------------------------------------- ;

// if you change Encode(...) / Decode(...), don't forget to update ComputeNumBytesNeeded(...)

bool BitStuffer2::EncodeSimple(Byte** ppByte, const vector<unsigned int>& dataVec) const
{
  if (!ppByte || dataVec.empty())
    return false;

  unsigned int maxElem = *max_element(dataVec.begin(), dataVec.end());
  int numBits = 0;
  while ((numBits < 32) && (maxElem >> numBits))
    numBits++;

  if (numBits >= 32)
    return false;

  Byte numBitsByte = (Byte)numBits;
  unsigned int numElements = (unsigned int)dataVec.size();
  unsigned int numUInts = (numElements * numBits + 31) / 32;

  // use the upper 2 bits to encode the type used for numElements: Byte, ushort, or uint
  int nb = NumBytesUInt(numElements);
  int bits67 = (nb == 4) ? 0 : 3 - nb;
  numBitsByte |= bits67 << 6;

  // bit5 = 0 means simple mode

  **ppByte = numBitsByte;
  (*ppByte)++;

  if (!EncodeUInt(ppByte, numElements, nb))
    return false;

  if (numUInts > 0)    // numBits can be 0, then we only write the header
    BitStuff(ppByte, dataVec, numBits);

  return true;
}

// -------------------------------------------------------------------------- ;

bool BitStuffer2::EncodeLut(Byte** ppByte, const vector<pair<unsigned int, unsigned int> >& sortedDataVec) const
{
  if (!ppByte || sortedDataVec.empty())
    return false;

  if (sortedDataVec[0].first != 0)    // corresponds to min
    return false;

  // collect the different values for the lut
  unsigned int numElem = (unsigned int)sortedDataVec.size();
  unsigned int indexLut = 0;

  m_tmpLutVec.resize(0);    // omit the 0 throughout that corresponds to min
  m_tmpIndexVec.resize(numElem);
  memset(&m_tmpIndexVec[0], 0, numElem * sizeof(unsigned int));

  for (unsigned int i = 1; i < numElem; i++)
  {
    unsigned int prev = sortedDataVec[i - 1].first;
    m_tmpIndexVec[sortedDataVec[i - 1].second] = indexLut;

    if (sortedDataVec[i].first != prev)
    {
      m_tmpLutVec.push_back(sortedDataVec[i].first);
      indexLut++;
    }
  }
  m_tmpIndexVec[sortedDataVec[numElem - 1].second] = indexLut;    // don't forget the last one

  // write first 2 data elements same as simple, but bit5 set to 1
  unsigned int maxElem = m_tmpLutVec.back();
  int numBits = 0;
  while ((numBits < 32) && (maxElem >> numBits))
    numBits++;

  if (numBits >= 32)
    return false;

  Byte numBitsByte = (Byte)numBits;

  // use the upper 2 bits to encode the type used for numElem: byte, ushort, or uint
  int n = NumBytesUInt(numElem);
  int bits67 = (n == 4) ? 0 : 3 - n;
  numBitsByte |= bits67 << 6;

  numBitsByte |= (1 << 5);    // bit 5 = 1 means lut mode

  **ppByte = numBitsByte;
  (*ppByte)++;

  if (!EncodeUInt(ppByte, numElem, n))    // numElements = numIndexes to lut
    return false;

  unsigned int nLut = (unsigned int)m_tmpLutVec.size();
  if (nLut < 1 || nLut >= 255)
    return false;

  **ppByte = (Byte)nLut + 1;    // size of lut, incl the 0
  (*ppByte)++;

  BitStuff(ppByte, m_tmpLutVec, numBits);    // lut

  int nBitsLut = 0;
  while (nLut >> nBitsLut)    // indexes are in [0 .. nLut]
    nBitsLut++;

  BitStuff(ppByte, m_tmpIndexVec, nBitsLut);    // indexes

  return true;
}

// -------------------------------------------------------------------------- ;

// if you change Encode(...) / Decode(...), don't forget to update ComputeNumBytesNeeded(...)

bool BitStuffer2::Decode(const Byte** ppByte, vector<unsigned int>& dataVec, int lerc2Version) const
{
  if (!ppByte)
    return false;

  Byte numBitsByte = **ppByte;
  (*ppByte)++;

  int bits67 = numBitsByte >> 6;
  int nb = (bits67 == 0) ? 4 : 3 - bits67;

  bool doLut = (numBitsByte & (1 << 5)) ? true : false;    // bit 5
  numBitsByte &= 31;    // bits 0-4;
  int numBits = numBitsByte;

  unsigned int numElements = 0;
  if (!DecodeUInt(ppByte, numElements, nb))
    return false;

  if (!doLut)
  {
    if (numBits > 0)    // numBits can be 0
    {
      if (lerc2Version >= 3)
        BitUnStuff(ppByte, dataVec, numElements, numBits);
      else
        BitUnStuff_Before_Lerc2v3(ppByte, dataVec, numElements, numBits);
    }
    else    // numBits = 0, all elements = 0
    {
      dataVec.resize(numElements);
      memset(&dataVec[0], 0, numElements * sizeof(unsigned int));
    }
  }
  else
  {
    Byte nLutByte = **ppByte;
    (*ppByte)++;

    int nLut = nLutByte - 1;

    // unstuff lut w/o the 0
    if (lerc2Version >= 3)
      BitUnStuff(ppByte, m_tmpLutVec, nLut, numBits);
    else
      BitUnStuff_Before_Lerc2v3(ppByte, m_tmpLutVec, nLut, numBits);

    int nBitsLut = 0;
    while (nLut >> nBitsLut)
      nBitsLut++;

    // unstuff indexes
    if (lerc2Version >= 3)
      BitUnStuff(ppByte, dataVec, numElements, nBitsLut);
    else
      BitUnStuff_Before_Lerc2v3(ppByte, dataVec, numElements, nBitsLut);

    // replace indexes by values
    m_tmpLutVec.insert(m_tmpLutVec.begin(), 0);    // put back in the 0
    for (unsigned int i = 0; i < numElements; i++)
      dataVec[i] = m_tmpLutVec[dataVec[i]];
  }

  return true;
}

// -------------------------------------------------------------------------- ;

unsigned int BitStuffer2::ComputeNumBytesNeededLut(const vector<pair<unsigned int, unsigned int> >& sortedDataVec, bool& doLut)
{
  unsigned int maxElem = sortedDataVec.back().first;
  unsigned int numElem = (unsigned int)sortedDataVec.size();

  int numBits = 0;
  while ((numBits < 32) && (maxElem >> numBits))
    numBits++;
  unsigned int numBytes = 1 + NumBytesUInt(numElem) + ((numElem * numBits + 7) >> 3);

  // go through and count how often the value changes
  int nLut = 0;
  for (unsigned int i = 1; i < numElem; i++)
    if (sortedDataVec[i].first != sortedDataVec[i - 1].first)
      nLut++;

  int nBitsLut = 0;
  while (nLut >> nBitsLut)
    nBitsLut++;

  unsigned int numBitsTotalLut = nLut * numBits;    // num bits w/o the 0
  unsigned int numBytesLut = 1 + NumBytesUInt(numElem) + 1 + ((numBitsTotalLut + 7) >> 3) + ((numElem * nBitsLut + 7) >> 3);

  doLut = numBytesLut < numBytes;
  return min(numBytesLut, numBytes);
}

// -------------------------------------------------------------------------- ;
// -------------------------------------------------------------------------- ;

void BitStuffer2::BitUnStuff_Before_Lerc2v3(const Byte** ppByte, vector<unsigned int>& dataVec,
                             unsigned int numElements, int numBits) const
{
  dataVec.resize(numElements, 0);    // init with 0

  unsigned int numUInts = (numElements * numBits + 31) / 32;
  unsigned int numBytes = numUInts * sizeof(unsigned int);
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wcast-qual"
#pragma GCC diagnostic ignored "-Wcast-align"
  unsigned int* arr = (unsigned int*)(*ppByte);
#pragma GCC diagnostic pop

  unsigned int* srcPtr = arr;
  srcPtr += numUInts;

  // needed to save the 0-3 bytes not used in the last UInt
  srcPtr--;
  unsigned int lastUInt = *srcPtr;
  unsigned int numBytesNotNeeded = NumTailBytesNotNeeded(numElements, numBits);
  unsigned int n = numBytesNotNeeded;

  while (n--)
  {
    unsigned int val;
    memcpy(&val, srcPtr, sizeof(unsigned int));
    val <<= 8;
    memcpy(srcPtr, &val, sizeof(unsigned int));
  }

  // do the un-stuffing
  srcPtr = arr;
  unsigned int* dstPtr = &dataVec[0];
  int bitPos = 0;

  for (unsigned int i = 0; i < numElements; i++)
  {
    if (32 - bitPos >= numBits)
    {
      unsigned int val;
      memcpy(&val, srcPtr, sizeof(unsigned int));
      unsigned int n = val << bitPos;

      *dstPtr++ = n >> (32 - numBits);
      bitPos += numBits;
      if (bitPos == 32)    // shift >= 32 is undefined
      {
        bitPos = 0;
        srcPtr++;
      }
    }
    else
    {
      unsigned int val;
      memcpy(&val, srcPtr, sizeof(unsigned int));
      srcPtr++;
      unsigned int n = val << bitPos;
      *dstPtr = n >> (32 - numBits);
      bitPos -= (32 - numBits);
      memcpy(&val, srcPtr, sizeof(unsigned int));
      *dstPtr++ |= val >> (32 - bitPos);
    }
  }

  if (numBytesNotNeeded > 0)
    memcpy(srcPtr, &lastUInt, sizeof(unsigned int));  // restore the last UInt

  *ppByte += numBytes - numBytesNotNeeded;
}

// -------------------------------------------------------------------------- ;

// starting with version Lerc2v3: integer >> into local uint buffer, plus final memcpy

void BitStuffer2::BitStuff(Byte** ppByte, const vector<unsigned int>& dataVec, int numBits) const
{
  unsigned int numElements = (unsigned int)dataVec.size();
  unsigned int numUInts = (numElements * numBits + 31) / 32;
  unsigned int numBytes = numUInts * sizeof(unsigned int);

  m_tmpBitStuffVec.resize(numUInts);
  unsigned int* dstPtr = &m_tmpBitStuffVec[0];

  memset(dstPtr, 0, numBytes);

  // do the stuffing
  const unsigned int* srcPtr = &dataVec[0];
  int bitPos = 0;

  for (unsigned int i = 0; i < numElements; i++)
  {
    if (32 - bitPos >= numBits)
    {
      *dstPtr |= (*srcPtr++) << bitPos;
      bitPos += numBits;
      if (bitPos == 32)    // shift >= 32 is undefined
      {
        dstPtr++;
        bitPos = 0;
      }
    }
    else
    {
      *dstPtr++ |= (*srcPtr) << bitPos;
      *dstPtr |= (*srcPtr++) >> (32 - bitPos);
      bitPos += numBits - 32;
    }
  }

  // copy the bytes to the outgoing byte stream
  int numBytesUsed = numBytes - NumTailBytesNotNeeded(numElements, numBits);
  memcpy(*ppByte, &m_tmpBitStuffVec[0], numBytesUsed);

  *ppByte += numBytesUsed;
}

// -------------------------------------------------------------------------- ;

void BitStuffer2::BitUnStuff(const Byte** ppByte, vector<unsigned int>& dataVec,
  unsigned int numElements, int numBits) const
{
  dataVec.resize(numElements);

  unsigned int numUInts = (numElements * numBits + 31) / 32;
  unsigned int numBytes = numUInts * sizeof(unsigned int);

  m_tmpBitStuffVec.resize(numUInts);
  m_tmpBitStuffVec[numUInts - 1] = 0;    // set last uint to 0

  // copy the bytes from the incoming byte stream
  int numBytesUsed = numBytes - NumTailBytesNotNeeded(numElements, numBits);
  memcpy(&m_tmpBitStuffVec[0], *ppByte, numBytesUsed);

  // do the un-stuffing
  unsigned int* srcPtr = &m_tmpBitStuffVec[0];
  unsigned int* dstPtr = &dataVec[0];
  int bitPos = 0;
  int nb = 32 - numBits;

  for (unsigned int i = 0; i < numElements; i++)
  {
    if (nb - bitPos >= 0)
    {
      *dstPtr++ = ((*srcPtr) << (nb - bitPos)) >> nb;
      bitPos += numBits;
      if (bitPos == 32)    // shift >= 32 is undefined
      {
        srcPtr++;
        bitPos = 0;
      }
    }
    else
    {
      *dstPtr = (*srcPtr++) >> bitPos;
      *dstPtr++ |= ((*srcPtr) << (64 - numBits - bitPos)) >> nb;
      bitPos -= nb;
    }
  }

  *ppByte += numBytesUsed;
}

// -------------------------------------------------------------------------- ;

