Merge branch 'new_psy' into at3plus-devat3plus-dev

7 files changed, 243 insertions, 31 deletions

diff --git a/src/atrac/atrac3_bitstream.cpp b/src/atrac/atrac3_bitstream.cpp
index 9bada6f..000eb92 100644
--- a/src/atrac/atrac3_bitstream.cpp
+++ b/src/atrac/atrac3_bitstream.cpp
@@ -34,12 +34,26 @@ using std::vector;
 using std::memset;
 
 static const uint32_t FixedBitAllocTable[TAtrac3Data::MaxBfus] = {
-  5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-  3, 3, 3, 3, 3, 3, 3, 3,
-  3, 2, 2, 1,
-  1, 0
+  4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
+  2, 2, 2, 2, 2, 1, 1, 1,
+  1, 1, 1, 1,
+  0, 0
 };
 
+#define EAQ 1
+
+#ifdef EAQ
+
+static constexpr size_t LOSY_NAQ_START = 18;
+static constexpr size_t BOOST_NAQ_END = 10;
+
+#else
+
+static constexpr size_t LOSY_NAQ_START = 31;
+static constexpr size_t BOOST_NAQ_END = 0;
+
+#endif
+
 std::vector<float> TAtrac3BitStreamWriter::ATH;
 TAtrac3BitStreamWriter::TAtrac3BitStreamWriter(ICompressedOutput* container, const TContainerParams& params, uint32_t bfuIdxConst)
     : Container(container)
@@ -120,34 +134,15 @@ uint32_t TAtrac3BitStreamWriter::VLCEnc(const uint32_t selector, const int manti
     return bitsUsed;
 }
 
-static inline int ToInt(float x) {
-#if defined(_MSC_VER) && !defined(_WIN64)
-    int n;
-    __asm {
-        fld x
-        fistp n
-    }
-    return n;
-#else
-    return lrint(x);
-#endif
-}
-
-static inline void CalcMantisas(const float* values, const uint32_t first, const uint32_t last, const float mul, int* mantisas) {
-    for (uint32_t j = 0, f = first; f < last; f++, j++) {
-        mantisas[f] = ToInt(values[j] * mul);
-    }
-}
-
 std::pair<uint8_t, uint32_t> TAtrac3BitStreamWriter::CalcSpecsBitsConsumption(const TSingleChannelElement& sce,
-                                                        const vector<uint32_t>& precisionPerEachBlocks, int* mantisas)
+    const vector<uint32_t>& precisionPerEachBlocks, int* mantisas, vector<float>& energyErr)
 {
 
     const vector<TScaledBlock>& scaledBlocks = sce.ScaledBlocks;
     const uint32_t numBlocks = precisionPerEachBlocks.size();
     uint32_t bitsUsed = numBlocks * 3;
 
-    auto lambda = [this, numBlocks, mantisas, &precisionPerEachBlocks, &scaledBlocks](bool clcMode, bool calcMant) {
+    auto lambda = [this, numBlocks, mantisas, &precisionPerEachBlocks, &scaledBlocks, &energyErr](bool clcMode, bool calcMant) {
         uint32_t bits = 0;
         for (uint32_t i = 0; i < numBlocks; ++i) {
             if (precisionPerEachBlocks[i] == 0)
@@ -159,7 +154,7 @@ std::pair<uint8_t, uint32_t> TAtrac3BitStreamWriter::CalcSpecsBitsConsumption(co
             const float mul = TAtrac3Data::MaxQuant[std::min(precisionPerEachBlocks[i], (uint32_t)7)];
             if (calcMant) {
                 const float* values = scaledBlocks[i].Values.data();
-                CalcMantisas(values, first, last, mul, mantisas);
+                energyErr[i] = QuantMantisas(values, first, last, mul, i > LOSY_NAQ_START, mantisas);
             }
             bits += clcMode ? CLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr) :
                 VLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr);
@@ -189,6 +184,25 @@ static inline bool CheckBfus(uint16_t* numBfu, const vector<uint32_t>& precision
 
 static const std::pair<uint8_t, vector<uint32_t>> DUMMY_ALLOC{1, vector<uint32_t>{0}};
 
+bool ConsiderEnergyErr(const vector<float>& err, vector<uint32_t>& bits)
+{
+    if (err.size() < bits.size())
+        abort();
+
+    bool adjusted = false;
+    size_t lim = std::min((size_t)BOOST_NAQ_END, bits.size());
+    for (size_t i = 0; i < lim; i++) {
+        float e = err[i];
+        if (((e > 0 && e < 0.7) || e > 1.2) & (bits[i] < 7)) {
+            //std::cerr << "adjust: " << i << " e: " << e << " b: " << bits[i] << std::endl;
+            bits[i]++;
+            adjusted = true;
+        }
+    }
+
+    return adjusted;
+}
+
 std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(const TSingleChannelElement& sce,
     const uint16_t targetBits, int mt[TAtrac3Data::MaxSpecs], float laudness)
 {
@@ -211,6 +225,7 @@ std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(co
     }
 
     vector<uint32_t> precisionPerEachBlocks(numBfu);
+    vector<float> energyErr(numBfu);
     uint8_t mode;
     bool cont = true;
     while (cont) {
@@ -219,11 +234,17 @@ std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(co
         double minShift = -8;
         for (;;) {
             double shift = (maxShift + minShift) / 2;
-            const vector<uint32_t>& tmpAlloc = CalcBitsAllocation(scaledBlocks, numBfu, spread, shift, laudness);
-            auto consumption = CalcSpecsBitsConsumption(sce, tmpAlloc, mt);
+            vector<uint32_t> tmpAlloc = CalcBitsAllocation(scaledBlocks, numBfu, spread, shift, laudness);
+            energyErr.clear();
+            energyErr.resize(numBfu);
+            std::pair<uint8_t, uint32_t> consumption;
+
+            do {
+                consumption = CalcSpecsBitsConsumption(sce, tmpAlloc, mt, energyErr);
+            } while (ConsiderEnergyErr(energyErr, tmpAlloc));
 
             auto bitsUsedByTonal = EncodeTonalComponents(sce, tmpAlloc, nullptr);
-//            std::cerr << consumption.second << " |tonal: " << bitsUsedByTonal << " target: " << targetBits << " shift " << shift << " max | min " << maxShift << " " << minShift << " numBfu: " << numBfu << std::endl;
+            //std::cerr << consumption.second << " |tonal: " << bitsUsedByTonal << " target: " << targetBits << " shift " << shift << " max | min " << maxShift << " " << minShift << " numBfu: " << numBfu << std::endl;
             consumption.second += bitsUsedByTonal;
 
             if (consumption.second < targetBits) {
@@ -489,11 +510,25 @@ vector<uint32_t> TAtrac3BitStreamWriter::CalcBitsAllocation(const std::vector<TS
             bitsPerEachBlock[i] = 0;
         } else {
             const uint32_t fix = FixedBitAllocTable[i];
-            int tmp = spread * ( (float)scaledBlocks[i].ScaleFactorIndex/3.2) + (1.0 - spread) * fix - shift;
+            float x = 6;
+            if (i < 3) {
+                x = 2.8;
+            } else if (i < 10) {
+                x = 2.6;
+            } else if (i < 15) {
+                x = 3.3;
+            } else if (i <= 20) {
+                x = 3.6;
+            } else if (i <= 28) {
+                x = 4.2;
+            }
+            int tmp = spread * ( (float)scaledBlocks[i].ScaleFactorIndex / x) + (1.0 - spread) * fix - shift;
             if (tmp > 7) {
                 bitsPerEachBlock[i] = 7;
             } else if (tmp < 0) {
                 bitsPerEachBlock[i] = 0;
+            } else if (tmp == 0) {
+                bitsPerEachBlock[i] = 1;
             } else {
                 bitsPerEachBlock[i] = tmp;
             }
diff --git a/src/atrac/atrac3_bitstream.h b/src/atrac/atrac3_bitstream.h
index 93ee72e..34a6fc5 100644
--- a/src/atrac/atrac3_bitstream.h
+++ b/src/atrac/atrac3_bitstream.h
@@ -73,7 +73,7 @@ private:
 
     std::pair<uint8_t, uint32_t> CalcSpecsBitsConsumption(const TSingleChannelElement& sce,
                                                           const std::vector<uint32_t>& precisionPerEachBlocks,
-                                                          int* mantisas);
+                                                          int* mantisas, std::vector<float>& energyErr);
 
     void EncodeSpecs(const TSingleChannelElement& sce, NBitStream::TBitStream* bitStream,
                      const std::pair<uint8_t, std::vector<uint32_t>>&, const int mt[TAtrac3Data::MaxSpecs]);
diff --git a/src/atrac/atrac_scale.cpp b/src/atrac/atrac_scale.cpp
index d6f5960..a436aae 100644
--- a/src/atrac/atrac_scale.cpp
+++ b/src/atrac/atrac_scale.cpp
@@ -19,6 +19,7 @@
 #include "atrac_scale.h"
 #include "atrac1.h"
 #include "atrac3.h"
+#include "util.h"
 #include <cmath>
 #include <iostream>
 #include <algorithm>
@@ -35,6 +36,100 @@ using std::abs;
 
 static const float MAX_SCALE = 1.0;
 
+float QuantMantisas(const float* in, const uint32_t first, const uint32_t last, const float mul, bool ea, int* const mantisas)
+{
+    float e1 = 0.0;
+    float e2 = 0.0;
+
+    const float inv2 = 1.0 / (mul * mul);
+
+    if (!ea) {
+        for (uint32_t j = 0, f = first; f < last; f++, j++) {
+            float t = in[j] * mul;
+            e1 += in[j] * in[j];
+            mantisas[f] = ToInt(t);
+            e2 += mantisas[f] * mantisas[f] * inv2;
+        }
+        return e1 / e2;
+    }
+
+    std::vector<std::pair<float, int>> candidates;
+    candidates.reserve(last - first);
+
+    for (uint32_t j = 0, f = first; f < last; f++, j++) {
+        float t = in[j] * mul;
+        e1 += in[j] * in[j];
+        mantisas[f] = ToInt(t);
+        e2 += mantisas[f] * mantisas[f] * inv2;
+
+        float delta = t - (std::truncf(t) + 0.5f);
+        // 0 ... 0.25 ... 0.5 ... 0.75 ... 1
+        //        ^----------------^ candidates to be rounded to opposite side
+        // to decrease overall energy error in the band
+        if (std::abs(delta) < 0.25f) {
+            candidates.push_back({delta, f});
+        }
+    }
+
+    if (candidates.empty()) {
+        return e1 / e2;
+    }
+
+    static auto cmp = [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
+        return std::abs(a.first) < std::abs(b.first);
+    };
+
+    std::sort(candidates.begin(), candidates.end(), cmp);
+
+    if (e2 < e1) {
+        for (const auto& x : candidates) {
+            auto f = x.second;
+            auto j = f - first;
+            auto t = in[j] * mul;
+            if (static_cast<float>(std::abs(mantisas[f])) < std::abs(t) && static_cast<float>(std::abs(mantisas[f])) < (mul - 1)) {
+                int m = mantisas[f];
+                if (m > 0) m++;
+                if (m < 0) m--;
+                if (m == 0) m = t > 0 ? 1 : -1;
+                auto ex = e2;
+                ex -= mantisas[f] * mantisas[f] * inv2;
+                ex += m * m * inv2;
+                if (std::abs(ex - e1) < std::abs(e2 - e1)) {
+                    mantisas[f] = m;
+                    e2 = ex;
+                }
+            }
+        }
+        return e1 / e2;
+    }
+
+    if (e2 > e1) {
+        for (const auto& x : candidates) {
+            auto f = x.second;
+            auto j = f - first;
+            auto t = in[j] * mul;
+            if (static_cast<float>(std::abs(mantisas[f])) > std::abs(t)) {
+                auto m = mantisas[f];
+                if (m > 0) m--;
+                if (m < 0) m++;
+
+                auto ex = e2;
+                ex -= mantisas[f] * mantisas[f] * inv2;
+                ex += m * m * inv2;
+
+                if (std::abs(ex - e1) < std::abs(e2 - e1)) {
+                    mantisas[f] = m;
+                    e2 = ex;
+                }
+            }
+        }
+        return e1 / e2;
+    }
+    return e1 / e2;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
 template<class TBaseData>
 TScaler<TBaseData>::TScaler() {
     for (int i = 0; i < 64; i++) {
diff --git a/src/atrac/atrac_scale.h b/src/atrac/atrac_scale.h
index 059360d..9a812f0 100644
--- a/src/atrac/atrac_scale.h
+++ b/src/atrac/atrac_scale.h
@@ -27,6 +27,8 @@
 
 namespace NAtracDEnc {
 
+float QuantMantisas(const float* in, uint32_t first, uint32_t last, float mul, bool ea, int* mantisas);
+
 struct TScaledBlock {
     TScaledBlock(uint8_t sfi) : ScaleFactorIndex(sfi) {}
     /* const */ uint8_t ScaleFactorIndex = 0;
diff --git a/src/atrac/atrac_scale_ut.cpp b/src/atrac/atrac_scale_ut.cpp
new file mode 100644
index 0000000..810c469
--- /dev/null
+++ b/src/atrac/atrac_scale_ut.cpp
@@ -0,0 +1,67 @@
+/*
+ * This file is part of AtracDEnc.
+ *
+ * AtracDEnc is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * AtracDEnc 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
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with AtracDEnc; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "atrac_scale.h"
+
+#include <gtest/gtest.h>
+
+using namespace NAtracDEnc;
+
+TEST(Quant, SaveEnergyLost) {
+    struct TTestData {
+        const std::vector<float> In;
+        const float S;
+        const float Q;
+        const float Diff;
+    };
+
+    const std::vector<TTestData> testData {
+        {{-2.35, -0.84,  0.65, -1.39,  1.25, -0.41, -0.85,  0.89}, 2.35001, 2.5, 0.5},
+        {{-1.26,  1.26, -1.26,  1.26, -1.26,  1.26, -1.26,  1.26}, 2.35001, 2.5, 0.4},
+        {{-0.32,  0.13,  0.28,  0.35,  0.63,  0.86,  0.63,  0.04}, 1, 15.5, 0.03}
+    };
+
+    for (const auto& td : testData) {
+        const std::vector<float>& in = td.In;
+        const float scale = td.S;
+
+        float e1 = 0.0;
+
+        std::vector<float> scaled;
+        scaled.reserve(in.size());
+
+        for (auto x : in) {
+            e1 += x * x;
+            scaled.push_back(x / scale);
+        }
+
+        std::vector<int> mantisas;
+        mantisas.resize(in.size());
+
+        QuantMantisas(scaled.data(), 0, mantisas.size(), td.Q, true, mantisas.data());
+
+        float e2 = 0.0;
+        for (auto x : mantisas) {
+            auto t = x * (scale / td.Q);
+            e2 += t * t;
+        }
+
+        EXPECT_TRUE(std::abs(e2 - e1) < td.Diff);
+        std::cerr << "(e1): " << e1 << " (e2): " << e2 << " (e2-e1) " << e2 - e1 << " (e2/e1) " << e2 / e1 << std::endl;
+    }
+}
diff --git a/src/util.h b/src/util.h
index cc3249b..ff4399e 100644
--- a/src/util.h
+++ b/src/util.h
@@ -90,3 +90,15 @@ inline T CalcEnergy(const std::vector<T>& in) {
         });
 }
 
+inline int ToInt(float x) {
+#if defined(_MSC_VER) && !defined(_WIN64)
+    int n;
+    __asm {
+        fld x
+        fistp n
+    }
+    return n;
+#else
+    return lrint(x);
+#endif
+}
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index f5230aa..809070e 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -15,6 +15,7 @@ set(atracdenc_ut
     ${CMAKE_SOURCE_DIR}/src/atracdenc_ut.cpp
     ${CMAKE_SOURCE_DIR}/src/atrac3denc_ut.cpp
     ${CMAKE_SOURCE_DIR}/src/transient_detector_ut.cpp
+    ${CMAKE_SOURCE_DIR}/src/atrac/atrac_scale_ut.cpp
 )
 
 add_executable(atracdenc_ut ${atracdenc_ut})