Use ATH during ATRAC1 encoding

4 files changed, 47 insertions, 15 deletions

diff --git a/src/atrac/atrac1_bitalloc.cpp b/src/atrac/atrac1_bitalloc.cpp
index 69485db..a2e63fa 100644
--- a/src/atrac/atrac1_bitalloc.cpp
+++ b/src/atrac/atrac1_bitalloc.cpp
@@ -98,6 +98,30 @@ uint32_t TBitsBooster::ApplyBoost(std::vector<uint32_t>* bitsPerEachBlock, uint3
     return surplus;
 }
 
+std::vector<TFloat> TAtrac1SimpleBitAlloc::ATHLong;
+
+TAtrac1SimpleBitAlloc::TAtrac1SimpleBitAlloc(ICompressedOutput* container, uint32_t bfuIdxConst, bool fastBfuNumSearch)
+    : TAtrac1BitStreamWriter(container)
+    , BfuIdxConst(bfuIdxConst)
+    , FastBfuNumSearch(fastBfuNumSearch)
+{
+    if (ATHLong.size()) {
+        return;
+    }
+    ATHLong.reserve(MaxBfus);
+    auto ATHSpec = CalcATH(512, 44100);
+    for (size_t bandNum = 0; bandNum < this->NumQMF; ++bandNum) {
+        for (size_t blockNum = this->BlocksPerBand[bandNum]; blockNum < this->BlocksPerBand[bandNum + 1]; ++blockNum) {
+           const size_t specNumStart =  this->SpecsStartLong[blockNum];
+           float x = 999;
+           for (size_t line = specNumStart; line < specNumStart + this->SpecsPerBlock[blockNum]; line++) {
+                x = fmin(x, ATHSpec[line]);
+           }
+           x = pow(10, 0.1 * x);
+           ATHLong.push_back(x);
+        }
+    }
+}
 
 vector<uint32_t> TAtrac1SimpleBitAlloc::CalcBitsAllocation(const std::vector<TScaledBlock>& scaledBlocks,
                                                            const uint32_t bfuNum,
@@ -106,17 +130,22 @@ vector<uint32_t> TAtrac1SimpleBitAlloc::CalcBitsAllocation(const std::vector<TSc
                                                            const TBlockSize& blockSize) {
     vector<uint32_t> bitsPerEachBlock(bfuNum);
     for (size_t i = 0; i < bitsPerEachBlock.size(); ++i) {
-        const uint32_t fix = blockSize.LogCount[BfuToBand(i)] ? FixedBitAllocTableShort[i] : FixedBitAllocTableLong[i];
-        int tmp = spread * ( (TFloat)scaledBlocks[i].ScaleFactorIndex/3.2) + (1.0 - spread) * fix - shift;
-        if (tmp > 16) {
-            bitsPerEachBlock[i] = 16;
-        } else if (tmp < 2) {
+        bool shortBlock = blockSize.LogCount[BfuToBand(i)];
+        const uint32_t fix = shortBlock ? FixedBitAllocTableShort[i] : FixedBitAllocTableLong[i];
+        if (!shortBlock && scaledBlocks[i].MaxEnergy < ATHLong[i]) {
             bitsPerEachBlock[i] = 0;
         } else {
-            bitsPerEachBlock[i] = tmp;
+            int tmp = spread * ( (TFloat)scaledBlocks[i].ScaleFactorIndex/3.2) + (1.0 - spread) * fix - shift;
+            if (tmp > 16) {
+                bitsPerEachBlock[i] = 16;
+            } else if (tmp < 2) {
+                bitsPerEachBlock[i] = 0;
+            } else {
+                bitsPerEachBlock[i] = tmp;
+            }
         }
     }
-    return bitsPerEachBlock;	
+    return bitsPerEachBlock;
 }
 
 uint32_t TAtrac1SimpleBitAlloc::GetMaxUsedBfuId(const vector<uint32_t>& bitsPerEachBlock) {
diff --git a/src/atrac/atrac1_bitalloc.h b/src/atrac/atrac1_bitalloc.h
index 411bece..b1f6b54 100644
--- a/src/atrac/atrac1_bitalloc.h
+++ b/src/atrac/atrac1_bitalloc.h
@@ -59,16 +59,14 @@ class TAtrac1SimpleBitAlloc : public TAtrac1BitStreamWriter, public TBitsBooster
                                              const TFloat spread, const TFloat shift, const TBlockSize& blockSize);
     const uint32_t BfuIdxConst;
     const bool FastBfuNumSearch;
+    static std::vector<TFloat> ATHLong;
+
     uint32_t GetMaxUsedBfuId(const std::vector<uint32_t>& bitsPerEachBlock);
     uint32_t CheckBfuUsage(bool* changed, uint32_t curBfuId, const std::vector<uint32_t>& bitsPerEachBlock);
 public:
-    explicit TAtrac1SimpleBitAlloc(ICompressedOutput* container, uint32_t bfuIdxConst, bool fastBfuNumSearch)
-        : TAtrac1BitStreamWriter(container)
-        , BfuIdxConst(bfuIdxConst)
-        , FastBfuNumSearch(fastBfuNumSearch)
-    {}
+    TAtrac1SimpleBitAlloc(ICompressedOutput* container, uint32_t bfuIdxConst, bool fastBfuNumSearch);
     ~TAtrac1SimpleBitAlloc() {};
-     uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize) override;
+    uint32_t Write(const std::vector<TScaledBlock>& scaledBlocks, const TBlockSize& blockSize) override;
 };
 
 } //namespace NAtrac1
diff --git a/src/atrac/atrac_scale.cpp b/src/atrac/atrac_scale.cpp
index f24f7b0..edc44a6 100644
--- a/src/atrac/atrac_scale.cpp
+++ b/src/atrac/atrac_scale.cpp
@@ -52,14 +52,18 @@ TScaledBlock TScaler<TBaseData>::Scale(const TFloat* in, uint16_t len) {
     const TFloat scaleFactor = scaleIter->first;
     const uint8_t scaleFactorIndex = scaleIter->second;
     TScaledBlock res(scaleFactorIndex);
+    TFloat maxEnergy = 0.0;
     for (uint16_t i = 0; i < len; ++i) {
         TFloat scaledValue = in[i] / scaleFactor;
+        TFloat energy = in[i] * in[i];
+        maxEnergy = std::max(maxEnergy, energy);
         if (abs(scaledValue) >= 1.0) {
             cerr << "got "<< scaledValue << " it is wrong scalling" << endl;
             scaledValue = (scaledValue > 0) ? 0.99999 : -0.99999;
         }
         res.Values.push_back(scaledValue);
-	}
+    }
+    res.MaxEnergy = maxEnergy;
     return res;
 }
 
diff --git a/src/atrac/atrac_scale.h b/src/atrac/atrac_scale.h
index 855010e..b590662 100644
--- a/src/atrac/atrac_scale.h
+++ b/src/atrac/atrac_scale.h
@@ -28,9 +28,10 @@
 namespace NAtracDEnc {
 
 struct TScaledBlock {
-	TScaledBlock(uint8_t sfi) : ScaleFactorIndex(sfi) {}
+    TScaledBlock(uint8_t sfi) : ScaleFactorIndex(sfi) {}
     /* const */ uint8_t ScaleFactorIndex = 0;
     std::vector<TFloat> Values;
+    TFloat MaxEnergy;
 };
 
 class TBlockSize;