/*
 * SpectralQuantization.cpp
 *
 * Copyright 2021 HIMSA II K/S - www.himsa.com. Represented by EHIMA -
 * www.ehima.com
 *
 * 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.
 */

#include "SpectralQuantization.hpp"

#include <algorithm>
#include <cmath>

#include "SpectralDataTables.hpp"

namespace Lc3Enc {

SpectralQuantization::SpectralQuantization(uint16_t NE_, uint8_t fs_ind_)
    : NE(NE_),
      fs_ind(fs_ind_),
      X_q(nullptr),
      lastnz(0),
      nbits_trunc(0),
      rateFlag(0),
      lsbMode(0),
      gg(0),
      lastnz_trunc(0),
      gg_ind(0),
      datapoints(nullptr),
      gg_off(0),
      gg_min(0),
      nbits_offset(0),
      nbits_spec(0),
      nbits_spec_adj(0),
      nbits_est(0),
      reset_offset_old(false),
      nbits_offset_old(0),
      nbits_spec_old(0),
      nbits_est_old(0)

{
  X_q = new int16_t[NE];
}

SpectralQuantization::~SpectralQuantization() { delete[] X_q; }

void SpectralQuantization::updateGlobalGainEstimationParameter(
    uint16_t nbits, uint16_t nbits_spec_local) {
  if (reset_offset_old) {
    nbits_offset = 0;
  } else {
    nbits_offset =
        0.8 * nbits_offset_old +
        0.2 * std::min(40.f, std::max(-40.f, nbits_offset_old + nbits_spec_old -
                                                 nbits_est_old));
  }

  nbits_spec_adj = static_cast<uint16_t>(nbits_spec_local + nbits_offset + 0.5);

  gg_off = -std::min(115, nbits / (10 * (fs_ind + 1))) - 105 - 5 * (fs_ind + 1);
}

void SpectralQuantization::computeSpectralEnergy(double* E, const double* X_f) {
  for (uint16_t k = 0; k < NE / 4; k++) {
    E[k] = pow(2, -31);
    for (uint8_t n = 0; n <= 3; n++) {
      E[k] += X_f[4 * k + n] * X_f[4 * k + n];
    }
    E[k] = 10 * log10(E[k]);
  }
}

void SpectralQuantization::globalGainEstimation(const double* E) {
  // converted pseudo-code from page 49/50  (d09r02_F2F)
  int16_t fac = 256;
  //𝑔𝑔𝑖𝑛𝑑 = 255;
  gg_ind = 255;
  for (uint8_t iter = 0; iter < 8; iter++) {
    fac >>= 1;
    //𝑔𝑔𝑖𝑛𝑑 -= fac;
    gg_ind -= fac;
    double tmp = 0;
    uint8_t iszero = 1;
    // for (i = 𝑁𝐸/4-1; i >= 0; i--)
    for (int8_t i = NE / 4 - 1; i >= 0; i--) {
      // if (E[i]*28/20 < (𝑔𝑔𝑖𝑛𝑑+𝑔𝑔𝑜𝑓𝑓))
      if (E[i] * 28 / 20.0 < (gg_ind + gg_off)) {
        if (iszero == 0) {
          tmp += 2.7 * 28 / 20.0;
        }
      } else {
        // if ((𝑔𝑔𝑖𝑛𝑑+𝑔𝑔𝑜𝑓𝑓) < E[i]*28/20 - 43*28/20)
        if ((gg_ind + gg_off) < E[i] * 28 / 20.0 - 43 * 28 / 20.0) {
          // tmp += 2*E[i]*28/20 – 2*(𝑔𝑔𝑖𝑛𝑑+𝑔𝑔𝑜𝑓𝑓) - 36*28/20;
          tmp += 2 * E[i] * 28 / 20.0 - 2 * (gg_ind + gg_off) - 36 * 28 / 20.0;
        } else {
          // tmp += E[i]*28/20 – (𝑔𝑔𝑖𝑛𝑑+𝑔𝑔𝑜𝑓𝑓) + 7*28/20;
          tmp += E[i] * 28 / 20.0 - (gg_ind + gg_off) + 7 * 28 / 20.0;
        }
        iszero = 0;
      }
    }
    // if (tmp > 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐′ *1.4*28/20 && iszero == 0)
    if ((tmp > nbits_spec_adj * 1.4 * 28 / 20.0) && (iszero == 0)) {
      //𝑔𝑔𝑖𝑛𝑑 += fac;
      gg_ind += fac;
    }
  }
}

bool SpectralQuantization::globalGainLimitation(const double* const X_f) {
  // -> not precisely clear where this limitation should occur
  // -> Is the following converted pseudo code meant?
  // -> Are the chosen datatypes appropriate?
  double X_f_max = 0;
  for (uint16_t n = 0; n < NE; n++) {
    X_f_max = std::max(X_f_max, fabs(X_f[n]));
  }
  if (X_f_max > 0) {
    gg_min = ceil(28 * log10(X_f_max / (32768 - 0.375))) - gg_off;
  } else {
    gg_min = 0;
  }
  bool reset_offset = false;
  // if (𝑔𝑔𝑖𝑛𝑑 < 𝑔𝑔𝑚𝑖𝑛 || 𝑋𝑓𝑚𝑎𝑥 == 0)
  if ((gg_ind < gg_min) || X_f_max == 0) {
    //𝑔𝑔𝑖𝑛𝑑 = 𝑔𝑔𝑚𝑖𝑛;
    //𝑟𝑒𝑠𝑒𝑡𝑜𝑓𝑓𝑠𝑒𝑡 = 1;
    gg_ind = gg_min;
    reset_offset = true;
  } else {
    //𝑟𝑒𝑠𝑒𝑡𝑜𝑓𝑓𝑠𝑒𝑡 = 0;
    reset_offset = false;
  }
  return reset_offset;
}

void SpectralQuantization::quantizeSpectrum(const double* const X_f) {
  gg = pow(10.f, (gg_ind + gg_off) / 28.0f);
  for (uint16_t n = 0; n < NE; n++) {
    if (X_f[n] >= 0) {
      X_q[n] = X_f[n] / gg + 0.375;
    } else {
      X_q[n] = X_f[n] / gg - 0.375;
    }
  }
}

void SpectralQuantization::computeBitConsumption(uint16_t nbits,
                                                 uint8_t& modeFlag) {
  // if (nbits > (160 + 𝑓𝑠𝑖𝑛𝑑 * 160))
  if (nbits > (160 + fs_ind * 160)) {
    rateFlag = 512;
  } else {
    rateFlag = 0;
  }
  // if (nbits >= (480 + 𝑓𝑠𝑖𝑛𝑑 * 160))
  if (nbits >= (480 + fs_ind * 160)) {
    modeFlag = 1;
  } else {
    modeFlag = 0;
  }

  // lastnz = 𝑁𝐸;
  lastnz = NE;
  // while (lastnz>2 && 𝑋𝑞[lastnz-1] == 0 && 𝑋𝑞[lastnz-2] == 0)
  while ((lastnz > 2) && (X_q[lastnz - 1] == 0) && (X_q[lastnz - 2] == 0)) {
    lastnz -= 2;
  }

  //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 = 0;
  //𝑛𝑏𝑖𝑡𝑠𝑡𝑟𝑢𝑛𝑐 = 0;
  //𝑛𝑏𝑖𝑡𝑠𝑙𝑠𝑏 = 0;
  uint32_t nbits_est_local = 0;
  uint32_t nbits_trunc_local = 0;
  nbits_lsb = 0;
  lastnz_trunc = 2;
  uint16_t c = 0;
  for (uint16_t n = 0; n < lastnz; n = n + 2) {
    uint16_t t = c + rateFlag;
    // if (n > 𝑁𝐸/2)
    if (n > NE / 2) {
      t += 256;
    }
    // a = abs(𝑋𝑞[n]);
    uint16_t a = abs(X_q[n]);
    uint16_t a_lsb = a;
    // b = abs(𝑋𝑞[n+1]);
    uint16_t b = abs(X_q[n + 1]);
    uint16_t b_lsb = b;
    uint16_t lev = 0;
    // while (max(a,b) >= 4)
    uint8_t pki;
    while (std::max(a, b) >= 4) {
      pki = ac_spec_lookup[t + lev * 1024];
      //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 += ac_spec_bits[pki][16];
      nbits_est_local += ac_spec_bits[pki][16];
      if ((lev == 0) && (modeFlag == 1)) {
        //𝑛𝑏𝑖𝑡𝑠𝑙𝑠𝑏 += 2;
        nbits_lsb += 2;
      } else {
        //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 += 2*2048;
        nbits_est_local += 2 * 2048;
      }
      a >>= 1;
      b >>= 1;
      lev = std::min(lev + 1, 3);
    }
    pki = ac_spec_lookup[t + lev * 1024];
    uint16_t sym = a + 4 * b;
    //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 += ac_spec_bits[pki][sym];
    nbits_est_local += ac_spec_bits[pki][sym];
    // a_lsb = abs(𝑋𝑞[n]); -> implemented earlier
    // b_lsb = abs(𝑋𝑞[n+1]); -> implemented earlier
    //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 += (min(a_lsb,1) + min(b_lsb,1)) * 2048;
    // nbits_est_local += (std::min(a_lsb,static_cast<uint16_t>(1))
    //              + std::min(b_lsb,static_cast<uint16_t>(1))) * 2048;
    // alternative implementation (more clear, more performant?)
    if (a_lsb > 0) nbits_est_local += 2048;
    if (b_lsb > 0) nbits_est_local += 2048;
    if ((lev > 0) && (modeFlag == 1)) {
      a_lsb >>= 1;
      b_lsb >>= 1;
      // if (a_lsb == 0 && 𝑋𝑞[n] != 0)
      if (a_lsb == 0 && X_q[n] != 0) {
        //𝑛𝑏𝑖𝑡𝑠𝑙𝑠𝑏++;
        nbits_lsb++;
      }
      // if (b_lsb == 0 && 𝑋𝑞[n+1] != 0)
      if (b_lsb == 0 && X_q[n + 1] != 0) {
        //𝑛𝑏𝑖𝑡𝑠𝑙𝑠𝑏++;
        nbits_lsb++;
      }
    }

    // if ((𝑋𝑞[n] != 0 || 𝑋𝑞[n+1] != 0) && (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 <= 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐*2048))
    // if (((X_q[n] != 0) || (X_q[n+1] != 0)) && (nbits_est_local <=
    // nbits_spec*2048))
    if (((X_q[n] != 0) || (X_q[n + 1] != 0)) &&
        (ceil(nbits_est_local / 2048.0) <= nbits_spec)) {
      lastnz_trunc = n + 2;
      //𝑛𝑏𝑖𝑡𝑠𝑡𝑟𝑢𝑛𝑐 = 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡;
      nbits_trunc_local = nbits_est_local;
    }
    if (lev <= 1) {
      t = 1 + (a + b) * (lev + 1);
    } else {
      t = 12 + lev;
    }
    c = (c & 15) * 16 + t;
  }
  //𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 = ceil(𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡/2048) + 𝑛𝑏𝑖𝑡𝑠𝑙𝑠𝑏;
  nbits_est = ceil(nbits_est_local / 2048.0) + nbits_lsb;
  //𝑛𝑏𝑖𝑡𝑠𝑡𝑟𝑢𝑛𝑐 = ceil(𝑛𝑏𝑖𝑡𝑠𝑡𝑟𝑢𝑛𝑐/2048);
  nbits_trunc = ceil(nbits_trunc_local / 2048.0);
}

bool SpectralQuantization::globalGainAdjustment() {
  uint16_t t1[5] = {80, 230, 380, 530, 680};
  uint16_t t2[5] = {500, 1025, 1550, 2075, 2600};
  uint16_t t3[5] = {850, 1700, 2550, 3400, 4250};

  double delta;
  double delta2;
  // if (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < t1[𝑓𝑠𝑖𝑛𝑑])
  if (nbits_est < t1[fs_ind]) {
    // delta = (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡+48)/16;
    delta = (nbits_est + 48) / static_cast<double>(16);
  }
  // else if (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < t2[𝑓𝑠𝑖𝑛𝑑])
  else if (nbits_est < t2[fs_ind]) {
    // tmp1 = t1[𝑓𝑠𝑖𝑛𝑑]/16+3;
    // tmp2 = t2[𝑓𝑠𝑖𝑛𝑑]/48;
    // delta = (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡-t1[𝑓𝑠𝑖𝑛𝑑])*(tmp2-tmp1)/(t2[𝑓𝑠𝑖𝑛𝑑]-t1[𝑓𝑠𝑖𝑛𝑑]) + tmp1;
    double tmp1 = t1[fs_ind] / static_cast<double>(16) + 3;
    double tmp2 = t2[fs_ind] / static_cast<double>(48);
    delta =
        (nbits_est - t1[fs_ind]) * (tmp2 - tmp1) / (t2[fs_ind] - t1[fs_ind]) +
        tmp1;
  }
  // else if (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < t3[𝑓𝑠𝑖𝑛𝑑])
  else if (nbits_est < t3[fs_ind]) {
    // delta = 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡/48;
    delta = nbits_est / static_cast<double>(48);
  } else {
    // delta = t3[𝑓𝑠𝑖𝑛𝑑]/48;
    delta = t3[fs_ind] / static_cast<double>(48);
  }
  // delta = nint(delta); // this looks like we need floating point
  // nint(.) is the rounding-to-nearest-integer function.
  delta = static_cast<int16_t>(delta +
                               0.5);  // this looks like we need floating point
  delta2 = delta + 2;

  uint16_t gg_ind_origin = gg_ind;

  /*
  if ((𝑔𝑔𝑖𝑛𝑑 < 255 && 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 > 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐) ||
      (𝑔𝑔𝑖𝑛𝑑 > 0 && 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐 – delta2))
  {
      if (𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐 – delta2)
      {
          𝑔𝑔𝑖𝑛𝑑 -= 1;
      }
      else if (𝑔𝑔𝑖𝑛𝑑 == 254 || 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 < 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐 + delta)
      {
          𝑔𝑔𝑖𝑛𝑑 += 1;
      }
      else
      {
          𝑔𝑔𝑖𝑛𝑑 += 2;
      }
          𝑔𝑔𝑖𝑛𝑑 = max(𝑔𝑔𝑖𝑛𝑑, 𝑔𝑔𝑚𝑖𝑛);
  }
  */
  if (((gg_ind < 255) && (nbits_est > nbits_spec)) ||
      ((gg_ind > 0) && (nbits_est < nbits_spec - delta2))) {
    if (nbits_est < nbits_spec - delta2) {
      gg_ind -= 1;
    } else if ((gg_ind == 254) || (nbits_est < nbits_spec + delta)) {
      gg_ind += 1;
    } else {
      gg_ind += 2;
    }
    gg_ind = std::max(gg_ind, gg_min);
  }

  return (gg_ind_origin != gg_ind);
}

void SpectralQuantization::run(const double* const X_f, uint16_t nbits,
                               uint16_t nbits_spec_local) {
  nbits_spec = nbits_spec_local;

  // 3.3.10.2 First global gain estimation  (d09r02_F2F)
  updateGlobalGainEstimationParameter(nbits, nbits_spec_local);

  double E[NE / 4];
  computeSpectralEnergy(E, X_f);
  globalGainEstimation(E);
  // Finally, the quantized gain index is limited such that
  // the quantized spectrum stays within the range [-32768,32767]
  bool reset_offset = globalGainLimitation(X_f);

  // 3.3.10.3 Quantization   (d09r02_F2F)
  quantizeSpectrum(X_f);

  // 3.3.10.4 Bit consumption    (d09r02_F2F)
  uint8_t modeFlag;
  computeBitConsumption(nbits, modeFlag);

  if (nullptr != datapoints) {
    datapoints->log("gg_ind", &gg_ind, sizeof(gg_ind));
    datapoints->log("gg", &gg, sizeof(gg));
    datapoints->log("X_q", X_q, sizeof(int16_t) * NE);
    datapoints->log("lastnz", &lastnz, sizeof(lastnz));
    datapoints->log("lastnz_trunc", &lastnz_trunc, sizeof(lastnz_trunc));
    datapoints->log("nbits_est", &nbits_est, sizeof(nbits_est));
    datapoints->log("nbits_trunc", &nbits_trunc, sizeof(nbits_trunc));
  }

  // 3.3.10.5 Truncation   (d09r02_F2F)
  for (uint16_t k = lastnz_trunc; k < lastnz; k++) {
    //𝑋𝑞[k] = 0;
    X_q[k] = 0;
  }

  // if (modeFlag == 1 && 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 > 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐)
  if ((modeFlag == 1) && (nbits_est > nbits_spec)) {
    lsbMode = 1;
  } else {
    lsbMode = 0;
  }

  if (nullptr != datapoints) {
    datapoints->log("lsbMode", &lsbMode, sizeof(lsbMode));
  }

  // Note: states needs to be transferred prior
  // to spectrum re-quantization!
  nbits_spec_old = nbits_spec;
  nbits_est_old = nbits_est;
  reset_offset_old = reset_offset;
  nbits_offset_old = nbits_offset;

  // 3.3.10.6 Global gain adjustment        (d09r02_F2F)
  bool isAdjusted = globalGainAdjustment();
  if (isAdjusted) {
    quantizeSpectrum(X_f);
    computeBitConsumption(nbits, modeFlag);

    // 3.3.10.5 Truncation   (d09r02_F2F)
    for (uint16_t k = lastnz_trunc; k < lastnz; k++) {
      //𝑋𝑞[k] = 0;
      X_q[k] = 0;
    }

    // if (modeFlag == 1 && 𝑛𝑏𝑖𝑡𝑠𝑒𝑠𝑡 > 𝑛𝑏𝑖𝑡𝑠𝑠𝑝𝑒𝑐)
    if (modeFlag == 1 && nbits_est > nbits_spec) {
      lsbMode = 1;
    } else {
      lsbMode = 0;
    }
  }
  if (nullptr != datapoints) {
    datapoints->log("isAdjusted", &isAdjusted, sizeof(isAdjusted));
    datapoints->log("gg_ind_adj", &gg_ind, sizeof(gg_ind));
    datapoints->log("gg_adj", &gg, sizeof(gg));
    datapoints->log("X_q_req", X_q, sizeof(int16_t) * NE);
    datapoints->log("lastnz_req", &lastnz_trunc, sizeof(lastnz_trunc));
    datapoints->log("nbits_est_req", &nbits_est, sizeof(nbits_est));
    datapoints->log("nbits_trunc_req", &nbits_trunc, sizeof(nbits_trunc));
    datapoints->log("lsbMode_req", &lsbMode, sizeof(lsbMode));
  }
}

void SpectralQuantization::registerDatapoints(DatapointContainer* datapoints_) {
  datapoints = datapoints_;
  if (nullptr != datapoints) {
    datapoints->addDatapoint("gg_off", &gg_off, sizeof(gg_off));
    datapoints->addDatapoint("gg_min", &gg_min, sizeof(gg_min));
    datapoints->addDatapoint("nbits_offset", &nbits_offset,
                             sizeof(nbits_offset));
  }
}

}  // namespace Lc3Enc