/// ハフマン木の節を表す構造体
template<typename T> struct HuffmanTreeNode
{
  /// 子の節の添字(節番号)に対する型
  typedef typename std::vector< HuffmanTreeNode<T> >::size_type size_type;

  T value;         /// 値
  size_t count;    /// 出現頻度
  size_type left;  /// 左の子の節番号
  size_type right; /// 右の子の節番号

  /* コンストラクタ */

  /// 全要素を指定して構築
  ///
  /// t 値
  /// c 出現頻度
  /// l,r 左右の子の節番号
  HuffmanTreeNode( T t, size_t c, size_type l, size_type r )
  : value( t ), count( c ), left( l ), right( r ) {}

  /* メンバ関数 */

  /// 節の大小関係を評価する
  ///
  /// 自分自身が引数の節よりも小さければ true を返す。
  /// 大小関係は出現頻度 count で評価する。
  ///
  /// dest 比較対象の節
  bool operator<( const HuffmanTreeNode& dest ) const
  { return( count < dest.count ); }
};

/*
  MakeHuffmanTree : forest を森としてハフマン木を作成する

  各節ははじめから子を持たない状態になっているものとする
  戻り値として根の位置(先頭要素からの距離)を返す
*/
template<typename T> typename std::vector< HuffmanTreeNode<T> >::difference_type
MakeHuffmanTree( typename std::vector< HuffmanTreeNode<T> >* forest )
{
  // 発生頻度によりソートする
  std::sort( forest->begin(), forest->end() );

  // 根の位置
  typename std::vector< HuffmanTreeNode<T> >::difference_type root = 0;
  // 木にする節の位置
  typename std::vector< HuffmanTreeNode<T> >::size_type sz = 0;

  while ( sz + 1 < forest->size() ) {
    // 木にする節の位置を示す反復子
    typename std::vector< HuffmanTreeNode<T> >::iterator s = forest->begin() + sz;
    // 木にする二つの節を連結する親の節(これが根の節となる)
    HuffmanTreeNode<T> node( T(), s->count + ( s + 1 )->count, sz, sz + 1 );
    // 挿入位置を決めて親の節を登録する
    typename std::vector< HuffmanTreeNode<T> >::iterator p = std::upper_bound( s + 2, forest->end(), node );
    root = std::distance( forest->begin(), p ); // insert前に根の節の位置を登録しておく
    forest->insert( p, node );
    sz += 2;
  }

  return( root );
}

typedef HuffmanTreeNode<unsigned char> UCharNode; // ハフマン木の要素
typedef std::vector<UCharNode> HuffmanTree;       // ハフマン木

/*
  CountPalet : 画像 draw 上のピクセルに対し、各パレットの出現頻度をカウントして palet に登録する

  RGB 各成分はまとめてカウンタへ登録する。
*/
void CountPalet( const DrawingArea_IF& draw, std::vector<size_t>* palet )
{
  // 出現頻度の初期化
  palet->assign( RGB::MAX + 1, 0 );

  Coord<int> c;
  RGB rgb;
  for ( c.y = 0 ; c.y < draw.size().y ; ++( c.y ) ) {
    for ( c.x = 0 ; c.x < draw.size().x ; ++( c.x ) ) {
      draw.point( c, &rgb );
      for ( RGB::size_type st = RGB::RED ; st < RGB::RGB_SIZE ; ++st )
        ++( ( *palet )[rgb[st]] );
    }
  }
}

/*
  MakeForest : パレットの出現頻度 counter から、ハフマン木作成用の森 forest を作成する
*/
void MakeForest( const std::vector<size_t>& counter, HuffmanTree* forest )
{
  forest->clear();

  for ( std::vector<size_t>::const_iterator cit = counter.begin() ;
        cit != counter.end() ; ++cit ) {
    if ( *cit == 0 ) continue;
    forest->push_back( UCharNode( std::distance( counter.begin(), cit ), *cit,
                                  forest->max_size(), forest->max_size() ) );
  }
}

/// ハフマン符号の型
typedef unsigned int HuffmanCode;

/*
  ハフマン符号型変数のビット構成

  ハフマン符号型変数の上位側にハフマン符号、下位側にその長さを保持する。
  ハフマン符号長を保持するビット列の長さは HUFFMAN_LENGTH_SIZE、
  ハフマン符号のビット長は HUFFMAN_CODE_SIZE で保持する。
  ハフマン符号は右詰めとする。
*/

/// ハフマン符号長を保持するビット列の長さ
const unsigned int HUFFMAN_LENGTH_SIZE = 5;
/// ハフマン符号を保持するビット列の長さ
const unsigned int HUFFMAN_CODE_SIZE = CHAR_BIT * sizeof( HuffmanCode ) / sizeof( char ) - HUFFMAN_LENGTH_SIZE;
/// ハフマン符号長を得るためのマスク
const HuffmanCode HUFFMAN_LENGTH_MASK = ~( ~0 << HUFFMAN_LENGTH_SIZE );

/*
  MakeHuffmanCodeFromTree : ハフマン木 huffmanTree の要素を符号化して配列表 huffmanCode に登録する

  処理は再帰的に行い、nodeIdx は次に処理する節の位置を表す。それまでに符号化した結果は code が持ち、
  level は node の高さ ( ハフマン符号長と等しい ) を表す。
*/
template<class T> bool MakeHuffmanCodeFromTree
( const typename std::vector< HuffmanTreeNode<T> >& huffmanTree,
  typename std::vector< HuffmanTreeNode<T> >::difference_type nodeIdx,
  HuffmanCode code, HuffmanCode level,
  std::map<T,HuffmanCode>* huffmanCode )
{
  if ( level > HUFFMAN_CODE_SIZE ) return( false );
  if ( static_cast<typename std::vector< HuffmanTreeNode<T> >::size_type>( nodeIdx ) >= huffmanTree.size() ) return( false );

  typename std::vector< HuffmanTreeNode<T> >::const_iterator node = huffmanTree.begin() + nodeIdx;

  if ( node->right != huffmanTree.max_size() )
    if ( ! MakeHuffmanCodeFromTree( huffmanTree, node->right, code * 2 + 1, level + 1, huffmanCode ) )
      return( false );
  if ( node->left != huffmanTree.max_size() )
    if ( ! MakeHuffmanCodeFromTree( huffmanTree, node->left, code * 2, level + 1, huffmanCode ) )
      return( false );
  if ( node->left == huffmanTree.max_size() && node->right == huffmanTree.max_size() )
    ( *huffmanCode )[node->value] = ( code << HUFFMAN_LENGTH_SIZE ) | level;

  return( true );
}

/*
  AddNodeToHuffmanTree : ハフマン符号をハフマン木 huffmanTree に追加する

  cit->first は値、cit->second は値に対するハフマン符号をそれぞれ表す。
*/
template<class T> void AddNodeToHuffmanTree
( typename std::map<T,HuffmanCode>::const_iterator cit,
  typename std::vector< HuffmanTreeNode<T> >* huffmanTree )
{
  HuffmanCode code = cit->second >> HUFFMAN_LENGTH_SIZE;
  HuffmanCode level = cit->second & HUFFMAN_LENGTH_MASK;

  HuffmanTreeNode<T> leaf( T(), 0, huffmanTree->max_size(), huffmanTree->max_size() );
  typename std::vector< HuffmanTreeNode<T> >::size_type nodeIdx = 0;

  for ( ; level > 0 ; --level ) {
    // 符号ビットが 1 なら右、0 なら左の子の節を選択
    if ( ( ( code >> ( level - 1 ) ) & 1 ) != 0 ) {
      if ( (*huffmanTree)[nodeIdx].right == huffmanTree->max_size() ) {
        huffmanTree->push_back( leaf );
        (*huffmanTree)[nodeIdx].right = huffmanTree->size() - 1;
      }
      nodeIdx = (*huffmanTree)[nodeIdx].right;
    } else {
      if ( (*huffmanTree)[nodeIdx].left == huffmanTree->max_size() ) {
        huffmanTree->push_back( leaf );
        (*huffmanTree)[nodeIdx].left = huffmanTree->size() - 1;
      }
      nodeIdx = (*huffmanTree)[nodeIdx].left;
    }
  }
  (*huffmanTree)[nodeIdx].value = cit->first;
}

/*
  MakeTreeFromHuffmanCode : ハフマン符号の配列表 huffmanCode からハフマン木 huffmanTree を作成する
*/
template<class T> void MakeTreeFromHuffmanCode
( const std::map<T,HuffmanCode>& huffmanCode,
  std::vector< HuffmanTreeNode<T> >* huffmanTree )
{
  HuffmanTreeNode<T> leaf( T(), 0, huffmanTree->max_size(), huffmanTree->max_size() );

  huffmanTree->clear();
  huffmanTree->push_back( leaf ); // 根の節を初期化
  for ( typename std::map<T,HuffmanCode>::const_iterator cit = huffmanCode.begin() ;
        cit != huffmanCode.end() ; ++cit )
    AddNodeToHuffmanTree( cit, huffmanTree );
}

typedef std::map<unsigned char,HuffmanCode> HuffmanCodeMap; // バイトデータとハフマンコードのマップ

/*
  WriteToArray : huffmanCode を encdata へ書き込む

  bitData は一時的にビット列を保持するキャッシュで、bitLen がバッファサイズをオー
  バーしたら encData へ出力する
*/
void WriteToArray( HuffmanCode huffmanCode, HuffmanCode* bitData, unsigned int* bitLen, std::vector<HuffmanCode>* encData )
{
  // HaffmanCode のビット数
  const unsigned int szHuffmanCode = sizeof( HuffmanCode ) * CHAR_BIT / sizeof( char );

  // huffmanCode をコードと長さに分解
  HuffmanCode code = huffmanCode >> HUFFMAN_LENGTH_SIZE;
  HuffmanCode length = huffmanCode & HUFFMAN_LENGTH_MASK;

  *bitLen += length;
  if ( *bitLen >= szHuffmanCode ) { // キャッシュがいっぱいになったら encdata へ出力
    *bitLen -= szHuffmanCode;
    encData->push_back( ( *bitData << ( length - *bitLen ) ) | ( code >> *bitLen ) );
    *bitData = code & ~( ~0 << *bitLen );
  } else {
    *bitData = ( *bitData << length ) | code;
  }
}

/*
  EncodeToHuffmanCode : 画像データ draw に対してハフマン符号化を行う

  パレットコードに対するハフマン符号列は huffmanCode にあり、符号化した結果は encData へ書き込まれる
  対象画像データが持つ全てのパレットコードが huffmanCode にあることを前提とする。そうでないデータが
  あった場合は処理を中断して false を返す
*/
bool EncodeToHuffmanCode
( const DrawingArea_IF& draw,
  const HuffmanCodeMap& huffmanCode,
  std::vector<HuffmanCode>* encData )
{
  if ( &draw == 0 || &huffmanCode == 0 || encData == 0 ) return( false );

  // HaffmanCode のビット数
  const unsigned int szHuffmanCode = sizeof( HuffmanCode ) * CHAR_BIT / sizeof( char );

  encData->clear();

  Coord<int> c;
  RGB rgb;
  HuffmanCode bitData = 0; // 符号化ビットのキャッシュ
  unsigned int bitLen = 0; // キャッシュしたビット列の長さ
  for ( c.y = 0 ; c.y < draw.size().y ; ++( c.y ) ) {
    for ( c.x = 0 ; c.x < draw.size().x ; ++( c.x ) ) {
      draw.point( c, &rgb );
      for ( RGB::size_type st = 0 ; st < RGB::RGB_SIZE ; ++st ) {
        HuffmanCodeMap::const_iterator cit = huffmanCode.find( rgb[st] );
        if ( cit == huffmanCode.end() ) return( false ); // 取得できないことはありえない
        WriteToArray( cit->second, &bitData, &bitLen, encData );
      }
    }
  }
  if ( bitLen > 0 )
    encData->push_back( bitData << ( szHuffmanCode - bitLen ) );

  return( true );
}

/*
  ReadFromTree : ハフマン符号列を cit で始まる反復子から読み取りながら、ハフマン木 huffmanTree をたどって値を rgb に取得する

  root はハフマン木の根を表し、encEnd はハフマン符号列のある配列の末尾の次を表す。
  bitData は cit から読み込んだビット列を保持するキャッシュで、bitLen が残りのビット長を保持する。
*/
bool ReadFromTree
( const HuffmanTree& huffmanTree,
  HuffmanTree::difference_type root,
  std::vector<HuffmanCode>::const_iterator* cit,
  std::vector<HuffmanCode>::const_iterator encEnd,
  RGB* rgb, HuffmanCode* bitData, unsigned int* bitLen )
{
  // HaffmanCode のビット数
  const unsigned int szHuffmanCode = sizeof( HuffmanCode ) * CHAR_BIT / sizeof( char );

  for ( RGB::size_type st = 0 ; st < RGB::RGB_SIZE ; ++st ) {
    HuffmanTree::const_iterator tree = huffmanTree.begin() + root; // たどる開始位置は root から
    while ( tree->left != huffmanTree.max_size() || tree->right != huffmanTree.max_size() ) {
      // キャッシュの再読み込み
      if ( *bitLen == 0 ) {
        if ( *cit == encEnd ) return( false );
        *bitData = *( ( *cit )++ );
        *bitLen = szHuffmanCode;
      }
      if ( ( ( *bitData >> ( *bitLen - 1 ) ) & 1 ) == 1 ) {
        if ( tree->right == huffmanTree.max_size() ) return( false ); // 右側がたどれない
        tree = huffmanTree.begin() + tree->right;
      } else {
        if ( tree->left == huffmanTree.max_size() ) return( false ); // 右側がたどれない
        tree = huffmanTree.begin() + tree->left;
      }
      --( *bitLen );
    }
    ( *rgb )[st] = tree->value;
  }

  return( true );
}

/*
  DecodeFromHuffmanCode : ハフマン符号列 encData から、ハフマン木 huffmanTree を使って複合化を行い、結果を画像領域 draw に描画する

  root は huffmanTree の根を表す。
*/
bool DecodeFromHuffmanCode
( const std::vector<HuffmanCode>& encData,
  const HuffmanTree& huffmanTree,
  HuffmanTree::difference_type root,
  DrawingArea_IF* draw )
{
  if ( &encData == 0 || &huffmanTree == 0 || draw == 0 ) return( false );

  Coord<int> c;
  RGB rgb;
  HuffmanCode bitData = 0; // encData から取り出したビット列を一時的に保持するキャッシュ
  unsigned int bitLen = 0; // キャッシュにある残りビット数
  std::vector<HuffmanCode>::const_iterator cit = encData.begin();
  for ( c.y = 0 ; c.y < ( draw->size() ).y ; ++( c.y ) ) {
    for ( c.x = 0 ; c.x < ( draw->size() ).x ; ++( c.x ) ) {
      if ( ! ReadFromTree( huffmanTree, root, &cit, encData.end(), &rgb, &bitData, &bitLen ) )
        return( false );
      draw->pset( c, rgb );
    }
  }

  return( true );
}

typedef std::pair<unsigned short int, unsigned short int> ChildNodeIndex;

/*
  TransHT_CopyChild : ハフマン木 srcTree から配列表 destTree への要素のコピー

  p は destTree 上にある親の節の left, right いずれかへのポインタで、そのコピー元が
  srcTree[childIndex] になる。
*/
void TransHT_CopyChild( const HuffmanTree& srcTree,
                        HuffmanTree::size_type childIndex,
                        std::vector<ChildNodeIndex>* destTree,
                        unsigned short int* p )
{
  const unsigned short int PADDING = 0x100; // 節番号に加える定数 ( 9 ビットめを 1 にする )

  if ( childIndex == srcTree.max_size() ) return;

  // 子の節が葉だった場合、色コードをセット
  if ( srcTree[childIndex].left == srcTree.max_size() &&
       srcTree[childIndex].right == srcTree.max_size() )
    *p = srcTree[childIndex].value;
  // そうでなかったら新たな要素を配列表に追加してその節番号をセット
  else {
    *p = destTree->size() + PADDING;
    TransHuffmanTree( srcTree, childIndex, destTree );
  }
}

/*
  TransHuffmanTree : ハフマン木 srcTree の srcIndex 番目の要素をそのまま配列表 destTree に収める

  処理の開始は srcIndex に根の節の番号を渡す。この時 destTree の先頭の要素は根の節となる
*/
void TransHuffmanTree( const HuffmanTree& srcTree,
                       HuffmanTree::size_type srcIndex,
                       std::vector<ChildNodeIndex>* destTree )
{
  if ( srcIndex >= srcTree.size() ) return;

  destTree->push_back( ChildNodeIndex( 0, 0 ) );
  std::vector<ChildNodeIndex>::size_type destIndex = destTree->size() - 1;

  // 左側の節をチェック
  TransHT_CopyChild( srcTree, srcTree[srcIndex].left, destTree, &( (*destTree)[destIndex].first ) );
  // 右側の節をチェック
  TransHT_CopyChild( srcTree, srcTree[srcIndex].right, destTree, &( (*destTree)[destIndex].second ) );
}

/*
  InverseHT_CheckChild : ハフマン木 destTree の子 child が節番号か RGB 成分かをチェックする
*/
void InverseHT_CheckChild( HuffmanTree* destTree, UCharNode::size_type* child )
{
  const unsigned short int PADDING = 0x100; // 節番号から減算する定数 ( 9 ビットめを 0 にする )

  // 節番号ならば 9 ビットめを 0 にする
  if ( *child > RGB::MAX ) {
    *child -= PADDING;
  // RGB 成分なら葉の節として追加・結合する
  } else {
    RGB::primary_type value = *child;
    *child = destTree->size();
    destTree->push_back( UCharNode( value, 0, destTree->max_size(), destTree->max_size() ) );
  }
}

/*
  InverseHuffmanTree : 配列表 srcTree をハフマン木 destTree に逆変換する

  srcTree の先頭の要素が根の節であることを前提とする。
*/
void InverseHuffmanTree( const std::vector<ChildNodeIndex>& srcTree, HuffmanTree* destTree )
{
  destTree->clear();

  // srcTree からそのまま destTree へ内容をコピー
  for ( std::vector<ChildNodeIndex>::const_iterator cit = srcTree.begin() ;
        cit != srcTree.end() ; ++cit )
    destTree->push_back( UCharNode( 0, 0, cit->first, cit->second ) );

  HuffmanTree::size_type sz = destTree->size();
  for ( HuffmanTree::size_type i = 0 ; i < sz ; ++i ) {
    // 左の節を逆変換
    InverseHT_CheckChild( destTree, &( (*destTree)[i].left ) );
    // 右の節を逆変換
    InverseHT_CheckChild( destTree, &( (*destTree)[i].right ) );
  }
}