#include <iostream>
#include <vector>
#include <string>
#include <queue>   // For a more efficient way to build Huffman tree
#include <climits> // For INT_MAX

using namespace std;

// 节点结构体
struct HNode
{
    int weight; // 权值
    int parent; // 双亲节点索引
    int Lchild; // 左孩子索引
    int Rchild; // 右孩子索引
};

// 编码结构体（用于哈夫曼编码表）
struct HCode
{
    char data;   // 字符数据
    string code; // 哈夫曼编码
};

class Huffman
{
private:
    HNode *HTree;      // 哈夫曼树节点数组
    HCode *HCodeTable; // 哈夫曼编码表
    int N;             // 叶子节点数量

    // 辅助函数：递归生成哈夫曼编码
    void generateCodeRecursive(int i, string currentCode);

public:
    // 构造函数
    Huffman() : HTree(nullptr), HCodeTable(nullptr), N(0) {}

    // 创建哈夫曼树
    // a[]: 字符的权重数组, n: 字符种类数, name[]: 字符数组
    void CreateHTree(int a[], int n, char name[]);

    // 创建哈夫曼编码表
    void CreateHCodeTable();

    // 编码 (将字符串s编码为字符串d)
    void Encode(const char *s, string &d);

    // 解码 (将字符串s解码为字符串d)
    void Decode(const string &s, string &d);

    // 打印编码表 (便于调试)
    void PrintCodeTable();

    // 析构函数，释放内存
    ~Huffman();
};

// 辅助函数：从当前未纳入树的节点中选择两个最小权值的节点
// 通过引用参数返回两个最小节点的索引
void SelectMin(int current_tree_size, HNode HTree[], int &idx1, int &idx2)
{
    long long min1 = LLONG_MAX; // 使用 long long 避免溢出 INT_MAX
    long long min2 = LLONG_MAX;
    idx1 = -1;
    idx2 = -1;

    for (int i = 0; i < current_tree_size; i++)
    {
        // 查找父节点为 -1 的节点（即未被合并的节点）
        if (HTree[i].parent == -1)
        {
            if (HTree[i].weight < min1)
            {
                // 更新第二个最小
                min2 = min1;
                idx2 = idx1;
                // 更新第一个最小
                min1 = HTree[i].weight;
                idx1 = i;
            }
            else if (HTree[i].weight < min2)
            { // 注意这里是 else if
                min2 = HTree[i].weight;
                idx2 = i;
            }
        }
    }
}

// 创建哈夫曼树
// a[]: 字符的权重数组, n: 字符种类数, name[]: 字符数组
void Huffman::CreateHTree(int a[], int n, char name[])
{
    if (n <= 1)
    { // 至少需要两个字符才能构建非平凡的哈夫曼树
        N = n;
        if (n == 1)
        { // 特殊情况，只有一个字符
            HTree = new HNode[1];
            HTree[0] = {a[0], -1, -1, -1};
            HCodeTable = new HCode[1];
            HCodeTable[0] = {name[0], "0"}; // 单个字符的编码为 "0" 或 ""
        }
        else
        { // n == 0
            HTree = nullptr;
            HCodeTable = nullptr;
        }
        return;
    }

    N = n;
    int totalNodes = 2 * n - 1; // 哈夫曼树总节点数
    HTree = new HNode[totalNodes];

    // 初始化叶子节点
    for (int i = 0; i < N; i++)
    {
        HTree[i].weight = a[i];
        HTree[i].parent = -1;
        HTree[i].Lchild = -1;
        HTree[i].Rchild = -1;
        // 如果需要，可以在这里初始化 HCodeTable 的 data 字段
        // 但通常 HCodeTable 的分配和初始化在 CreateHCodeTable 中完成
    }

    // 初始化非叶子节点（空的，待填充）
    for (int i = N; i < totalNodes; i++)
    {
        HTree[i] = {0, -1, -1, -1}; // 权重0，parent为-1，孩子为-1
    }

    // 构建哈夫曼树
    for (int i = N; i < totalNodes; i++)
    {
        int x, y;                  // x 和 y 将保存选中节点的索引
        SelectMin(i, HTree, x, y); // 在当前已有的 i 个节点中选择

        // 标记选择的节点已被合并
        HTree[x].parent = i;
        HTree[y].parent = i;

        // 设置新节点的左右孩子和权重
        HTree[i].Lchild = x;
        HTree[i].Rchild = y;
        HTree[i].weight = HTree[x].weight + HTree[y].weight;

        // 新节点 i 的 parent 仍为 -1，表示它尚未被合并到更高层的节点，
        // 并且它现在成为下一轮 SelectMin 的候选节点。
        // HTree[i].parent = -1; // 这一行是多余的，在循环开始时已经初始化为-1
    }

    // 初始化 HCodeTable 并填充字符数据
    HCodeTable = new HCode[N];
    for (int i = 0; i < N; ++i)
    {
        HCodeTable[i].data = name[i];
    }
}

// 辅助函数：递归生成哈夫曼编码
void Huffman::generateCodeRecursive(int i, string currentCode)
{
    // 检查是否是叶子节点
    if (HTree[i].Lchild == -1 && HTree[i].Rchild == -1)
    {
        // 找到叶子节点，将编码保存到 HCodeTable
        // 这里的 i 对应的是叶子节点在 HTree 数组中的索引
        // HCodeTable[k] 应该对应 HTree 中的第 k 个叶子节点
        // 由于我们在 CreateHTree 中将叶子节点放在 0 到 N-1 索引，所以这里直接用 i
        HCodeTable[i].code = currentCode;
        return;
    }

    // 递归左孩子，路径加 "0"
    if (HTree[i].Lchild != -1)
    {
        generateCodeRecursive(HTree[i].Lchild, currentCode + "0");
    }
    // 递归右孩子，路径加 "1"
    if (HTree[i].Rchild != -1)
    {
        generateCodeRecursive(HTree[i].Rchild, currentCode + "1");
    }
}

// 创建哈夫曼编码表
void Huffman::CreateHCodeTable()
{
    if (N == 0)
        return;
    if (N == 1)
    { // 只有一个字符的特殊处理
        // 编码已经在 CreateHTree 中设置
        return;
    }
    // 哈夫曼树的根节点通常是最后一个创建的节点，索引为 2*N - 2
    generateCodeRecursive(2 * N - 2, "");
}

// 编码 (将字符串s编码为字符串d)
void Huffman::Encode(const char *s, string &d)
{
    if (N == 0)
    {
        d = "";
        return;
    }
    d.clear(); // 清空目标字符串

    for (int j = 0; s[j] != '\0'; ++j)
    {
        char targetChar = s[j];
        bool found = false;
        for (int i = 0; i < N; ++i)
        {
            if (HCodeTable[i].data == targetChar)
            {
                d += HCodeTable[i].code;
                found = true;
                break;
            }
        }
        if (!found)
        {
            // 处理未找到的字符，例如抛出异常或输出错误
            cerr << "Error: Character '" << targetChar << "' not found in Huffman table." << endl;
            // 可以选择抛出异常或跳过该字符
        }
    }
}

// 解码 (将字符串s解码为字符串d)
void Huffman::Decode(const string &s, string &d)
{
    if (N == 0 || HTree == nullptr)
    {
        d = "";
        return;
    }
    if (N == 1)
    { // 只有一个字符的特殊处理
        char_t singleChar = HCodeTable[0].data;
        for (char_t codeChar : s)
        {
            if (codeChar == '0' || codeChar == '1')
            { // 假设单字符编码为 "0" 或 "1"
                d += singleChar;
            }
            else
            {
                cerr << "Error: Invalid code in single character Huffman tree for: " << codeChar << endl;
                return;
            }
        }
        return;
    }

    d.clear();                 // 清空目标字符串
    int rootIndex = 2 * N - 2; // 哈夫曼树的根节点索引
    int current = rootIndex;

    for (char bit : s)
    {
        if (bit == '0')
        {
            current = HTree[current].Lchild;
        }
        else if (bit == '1')
        {
            current = HTree[current].Rchild;
        }
        else
        {
            cerr << "Error: Invalid bit '" << bit << "' in encoded string." << endl;
            d.clear(); // 清空部分解码的结果
            return;
        }

        // 如果到达叶子节点
        if (HTree[current].Lchild == -1 && HTree[current].Rchild == -1)
        {
            // 找到该叶子节点对应的字符
            // 需要遍历 HCodeTable 找到 data，或者在 HNode 中存储 data_index
            // 鉴于目前 HCodeTable 的索引和 HTree 的叶子节点索引是一致的 (0~N-1)
            // 我们可以直接使用 current 作为 HCodeTable 的索引
            if (current >= 0 && current < N)
            {
                d += HCodeTable[current].data;
            }
            else
            {
                cerr << "Error: Decoded to an invalid leaf node index: " << current << endl;
                d.clear();
                return;
            }

            current = rootIndex; // 重置到根节点，继续解码下一个字符
        }
    }
}

// 打印编码表 (便于调试)
void Huffman::PrintCodeTable()
{
    if (N == 0)
    {
        cout << "Huffman table is empty." << endl;
        return;
    }
    cout << "Huffman Code Table:" << endl;
    for (int i = 0; i < N; ++i)
    {
        cout << "'" << HCodeTable[i].data << "': " << HCodeTable[i].code << endl;
    }
}

// 析构函数，释放内存
Huffman::~Huffman()
{
    delete[] HTree;
    delete[] HCodeTable;
    HTree = nullptr;
    HCodeTable = nullptr;
}

// =================================== 主函数示例 ===================================
int main()
{
    // 字符和它们的权重
    char_t chars[] = {'a', 'b', 'c', 'd', 'e', 'f'};
    int weights[] = {45, 13, 12, 16, 9, 5};
    int n = sizeof(chars) / sizeof(chars[0]);

    Huffman hf;
    hf.CreateHTree(weights, n, chars);
    hf.CreateHCodeTable();
    hf.PrintCodeTable();

    string encoded_str;
    const char *original_text = "abcdeff";
    hf.Encode(original_text, encoded_str);
    cout << "Original text: " << original_text << endl;
    cout << "Encoded text: " << encoded_str << endl;

    string decoded_str;
    hf.Decode(encoded_str, decoded_str);
    cout << "Decoded text: " << decoded_str << endl;

    cout << "\nTesting with single character:" << endl;
    char_t chars_single[] = {'Z'};
    int weights_single[] = {100};
    Huffman hf_single;
    hf_single.CreateHTree(weights_single, 1, chars_single);
    hf_single.CreateHCodeTable();
    hf_single.PrintCodeTable();
    string encoded_single;
    hf_single.Encode("Z", encoded_single);
    cout << "Encoded single: " << encoded_single << endl;
    string decoded_single;
    hf_single.Decode(encoded_single, decoded_single);
    cout << "Decoded single: " << decoded_single << endl;

    return 0;
}