/*

 * Copyright 2016 Google Inc.

 *

 * 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.

 */

// Entropy encoding (Huffman) utilities.

#include "guetzli/entropy_encode.h"

#include <assert.h>

#include <algorithm>

namespace guetzli {

bool SetDepth(int p0, HuffmanTree *pool, uint8_t *depth, int max_depth) {

  int stack[17];

  int level = 0;

  int p = p0;

  assert(max_depth <= 16);

  stack[0] = -1;

  while (true) {

    if (pool[p].index_left_ >= 0) {

      level++;

      if (level > max_depth) return false;

      stack[level] = pool[p].index_right_or_value_;

      p = pool[p].index_left_;

      continue;

    } else {

      depth[pool[p].index_right_or_value_] = static_cast<uint8_t>(level);

    }

    while (level >= 0 && stack[level] == -1) level--;

    if (level < 0) return true;

    p = stack[level];

    stack[level] = -1;

  }

}

// Sort the root nodes, least popular first.

static inline bool SortHuffmanTree(const HuffmanTree& v0,

                                   const HuffmanTree& v1) {

  if (v0.total_count_ != v1.total_count_) {

    return v0.total_count_ < v1.total_count_;

  }

  return v0.index_right_or_value_ > v1.index_right_or_value_;

}

// This function will create a Huffman tree.

//

// The catch here is that the tree cannot be arbitrarily deep.

// Brotli specifies a maximum depth of 15 bits for "code trees"

// and 7 bits for "code length code trees."

//

// count_limit is the value that is to be faked as the minimum value

// and this minimum value is raised until the tree matches the

// maximum length requirement.

//

// This algorithm is not of excellent performance for very long data blocks,

// especially when population counts are longer than 2**tree_limit, but

// we are not planning to use this with extremely long blocks.

//

// See http://en.wikipedia.org/wiki/Huffman_coding

void CreateHuffmanTree(const uint32_t *data,

                       const size_t length,

                       const int tree_limit,

                       HuffmanTree* tree,

                       uint8_t *depth) {

  // For block sizes below 64 kB, we never need to do a second iteration

  // of this loop. Probably all of our block sizes will be smaller than

  // that, so this loop is mostly of academic interest. If we actually

  // would need this, we would be better off with the Katajainen algorithm.

  for (uint32_t count_limit = 1; ; count_limit *= 2) {

    size_t n = 0;

    for (size_t i = length; i != 0;) {

      --i;

      if (data[i]) {

        const uint32_t count = std::max<uint32_t>(data[i], count_limit);

        tree[n++] = HuffmanTree(count, -1, static_cast<int16_t>(i));

      }

    }

    if (n == 1) {

      depth[tree[0].index_right_or_value_] = 1;      // Only one element.

      break;

    }

    std::sort(tree, tree + n, SortHuffmanTree);

    // The nodes are:

    // [0, n): the sorted leaf nodes that we start with.

    // [n]: we add a sentinel here.

    // [n + 1, 2n): new parent nodes are added here, starting from

    //              (n+1). These are naturally in ascending order.

    // [2n]: we add a sentinel at the end as well.

    // There will be (2n+1) elements at the end.

    const HuffmanTree sentinel(~static_cast<uint32_t>(0), -1, -1);

    tree[n] = sentinel;

    tree[n + 1] = sentinel;

    size_t i = 0;      // Points to the next leaf node.

    size_t j = n + 1;  // Points to the next non-leaf node.

    for (size_t k = n - 1; k != 0; --k) {

      size_t left, right;

      if (tree[i].total_count_ <= tree[j].total_count_) {

        left = i;

        ++i;

      } else {

        left = j;

        ++j;

      }

      if (tree[i].total_count_ <= tree[j].total_count_) {

        right = i;

        ++i;

      } else {

        right = j;

        ++j;

      }

      // The sentinel node becomes the parent node.

      size_t j_end = 2 * n - k;

      tree[j_end].total_count_ =

          tree[left].total_count_ + tree[right].total_count_;

      tree[j_end].index_left_ = static_cast<int16_t>(left);

      tree[j_end].index_right_or_value_ = static_cast<int16_t>(right);

      // Add back the last sentinel node.

      tree[j_end + 1] = sentinel;

    }

    if (SetDepth(static_cast<int>(2 * n - 1), &tree[0], depth, tree_limit)) {

      /* We need to pack the Huffman tree in tree_limit bits. If this was not

         successful, add fake entities to the lowest values and retry. */

      break;

    }

  }

}

}  // namespace guetzli