#include <iostream>
#include <string>
#include <vector>
#include <functional>
#include <unordered_map>
#include <queue>
using namespace std;
using uint = unsigned int;
class HuffmanTree {
public:
    // 這里的lambda表達式用來初始化function函數(shù)對象
    // priority_queue的構造函數(shù)指出，如果傳入一個參數(shù)，那這個參數(shù)用來初始化比較器對象
    // 如果傳入兩個參數(shù)，第一個是比較器對象，第二個是底層容器
    HuffmanTree()
        :min_heap_([](Node* n1, Node* n2)->bool {return n1->weight_ > n2->weight_; })
        , root_(nullptr)
    {}
    ~HuffmanTree() {
        init();
        cout << "已釋放所有內存！" << endl;
    }
    // 根據(jù)字符串創(chuàng)建哈夫曼樹
    void create(const string& str) {
        if (root_ != nullptr) {
            cout << "哈夫曼樹初始化..." << endl;
            init();
            cout << "初始化完成！" << endl;
        }
        // 統(tǒng)計頻率（權重）
        unordered_map<char, uint> w_map;
        for (char c : str) {
            w_map[c]++;
        }
        // 遍歷w_map，把所有的字符對應的權重放入小根堆，按照權重排序
        for (pair<const char, uint>& p : w_map) {
            min_heap_.push(new Node(p.first, p.second));
        }
        // 根據(jù)優(yōu)先級隊列，從小根堆中取出節(jié)點，構建哈夫曼樹
        while (min_heap_.size() > 1) {
            Node* n1 = min_heap_.top();
            min_heap_.pop();
            Node* n2 = min_heap_.top();
            min_heap_.pop();
            Node* node = new Node('\0', n1->weight_ + n2->weight_);  // 內部節(jié)點存\0
            node->left_ = n1;
            node->right_ = n2;
            min_heap_.push(node);
        }
        root_ = min_heap_.top();
        min_heap_.pop();
        // 創(chuàng)建完哈夫曼樹，直接對傳入的海量字符進行編碼并存儲到code_map_
        create_huffman_code(str);
    }
    string get_code(const string& str) {
        // 利用哈夫曼樹對str編碼并返回
        string code;
        for (char c : str) {
            code += code_map_[c];
        }
        return code;
    }
    void show_huffman_code() const {
        // 打印哈夫曼編碼
        for (const auto& pair : code_map_) {
            cout << pair.first << " : " << pair.second << endl;
        }
    }
    string decode(const string& encode_str) {
        Node* cur = root_;
        string decode_str;
        for (char c : encode_str) {
            if (c == '0') {
                cur = cur->left_;
            }
            else {
                cur = cur->right_;
            }
            if (cur->left_ == nullptr && cur->right_ == nullptr) {
                // 到達葉子節(jié)點
                decode_str.push_back(cur->data_);
                cur = root_;
            }
        }
        return decode_str;
    }
    uint get_wpl() {
        if (root_ == nullptr) {
            return 0;
        }
        if (root_->left_ == nullptr && root_->right_ == nullptr) {
            // 對于葉子節(jié)點，直接返回自己的weight * depth
            return root_->weight_ * 1;
        }
        else {
            // 對于內部節(jié)點，直接返回從子節(jié)點拿到的weight之和
            return get_w(root_->left_, 2) + get_w(root_->right_, 2);
        }
    }
private:
    struct Node {
        Node(char data, uint weight)
            :data_(data)
            , weight_(weight)
            , left_(nullptr)
            , right_(nullptr)
        {}
        char data_;
        uint weight_;
        Node* left_;
        Node* right_;
    };
private:
    // 防止當前對象重新構建哈夫曼樹，釋放所有的節(jié)點，然后初始化私有成員
    void init() {
        // 釋放哈夫曼樹的節(jié)點
        if (root_ != nullptr) {
            queue<Node*> q;
            q.push(root_);
            while (!q.empty()) {
                Node* node = q.front();
                q.pop();
                if (node->left_ != nullptr) {
                    q.push(node->left_);
                }
                if (node->right_ != nullptr) {
                    q.push(node->right_);
                }
                delete node;
            }
            MinHeap empty([](Node* n1, Node* n2)->bool {return n1->weight_ > n2->weight_; });
            swap(empty, min_heap_);
            code_map_.clear();
        }
    }
    void create_huffman_code(const string& str) {
        string code;
        create_huffman_code(root_, code);
    }
    void create_huffman_code(Node* node, string code) {
        if (node->left_ == nullptr && node->right_ == nullptr) {
            code_map_[node->data_] = code;
            return;
        }
        create_huffman_code(node->left_, code + "0");
        create_huffman_code(node->right_, code + "1");
    }
    uint get_w(Node* node, int depth) {
        if (node == nullptr) {
            return 0;
        }
        if (node->left_ == nullptr && node->right_ == nullptr) {
            // 對于葉子節(jié)點，直接返回自己的weight * depth
            return node->weight_ * depth;
        }
        else {
            // 對于內部節(jié)點，直接返回從子節(jié)點拿到的weight之和
            return get_w(node->left_, depth + 1) + get_w(node->right_, depth + 1);
        }
    }
private:    
    Node* root_;
    unordered_map<char, string> code_map_;  // 存儲字符對應的哈夫曼編碼
    using MinHeap = priority_queue<Node*, vector<Node*>, function<bool(Node*, Node*)>>;
    MinHeap min_heap_; // 構建哈夫曼樹的時候使用小根堆
};
int main() {
    string str = "Aa";
    HuffmanTree htree;
    htree.create(str);
    htree.show_huffman_code();
    cout << htree.get_wpl() << endl;
    str = "ABC";
    htree.create(str);
    htree.show_huffman_code();
    cout << htree.get_wpl() << endl;;
    string encode = htree.get_code(str);
    cout << "encode:" << encode << endl;
    cout << "decode:" << htree.decode(encode) << endl;
    return 0;
}