#include <iostream>
#include <fstream>
#include <string>
#include <bitset>
#include <vector>
#include <queue>
#include <unordered_map>
#include <random>

// 游程编码
std::string runLengthEncoding(const std::string& input) {
    std::string output;
    char prev = input[0];
    int count = 1;
    for (int i = 1; i < input.length(); i++) {
        if (input[i] != prev) {
            output += std::to_string(count) + prev;
            prev = input[i];
            count = 1;
        }
        else {
            count++;
        }
    }
    output += std::to_string(count) + prev;
    return output;
}

// Huffman 编码
struct Node {
    char symbol;
    int frequency;
    Node* left;
    Node* right;

    Node(char s, int f) : symbol(s), frequency(f), left(nullptr), right(nullptr) {}
    Node(int f) : symbol('\0'), frequency(f), left(nullptr), right(nullptr) {}
};

struct CompareNode {
    bool operator()(const Node* lhs, const Node* rhs) const {
        return lhs->frequency > rhs->frequency;
    }
};
// 从哈夫曼树中获取编码
void buildCode(Node* node, std::string code, std::unordered_map<char, std::string>& codes) {
    if (!node) {
        return;
    }
    if (node->symbol != '\0') {
        codes[node->symbol] = code;
        return;
    }
    buildCode(node->left, code + "0", codes);
    buildCode(node->right, code + "1", codes);
}

std::priority_queue<Node*, std::vector<Node*>, CompareNode> pq;

std::string huffmanEncoding(const std::string& input) {
    std::unordered_map<char, int> frequencyMap;
    for (char c : input) {
        frequencyMap[c]++;
    }
    // 优先队列
    for (auto freqPair : frequencyMap) {
        Node* node = new Node(freqPair.first, freqPair.second);
        pq.push(node);
    }

    // 构造哈夫曼树
    while (pq.size() > 1) {
        Node* left = pq.top();
        pq.pop();
        Node* right = pq.top();
        pq.pop();
        Node* parent = new Node(left->frequency + right->frequency);
        parent->left = left;
        parent->right = right;
        pq.push(parent);
    }
    // 保存哈夫曼编码
    std::unordered_map<char, std::string> codes;
    // 根据哈夫曼树构造哈夫曼编码
    buildCode(pq.top(), "", codes);

    std::string output;
    for (char c : input) {
        output += codes[c];
    }
    return output;
}

// 加密和解密编码(XOR操作是可逆的,所以加密解密是一个函数)
std::string xorEncryptDecrypt(const std::string& input, const std::string& key) {
    int keyLen = key.size();
    std::string output;
    for (int i = 0; i < input.length(); i++) {
        int x = input[i] - '0';
        int y = key[i % keyLen] - '0';
        int num = x ^ y;
        char num_char = '0' + x ^ y;
        output += num_char; // XOR
    }
    return output;
}

// 随机生成长度为n的key
std::string generateRandomKey(const int n) {
    std::string binaryString;
    std::random_device rd;  // 获取随机数种子
    std::mt19937 generator(rd()); // 标准的梅森旋转算法
    std::uniform_int_distribution<> distribution(0, 1); // 定义一个均匀分布从0到1

    for (size_t i = 0; i < n; ++i) {
        char bit = distribution(generator) + '0'; // 生成0或1,并转换为字符'0'或'1'
        binaryString.push_back(bit);
    }
    return binaryString;
}


// Huffman 译码
std::string huffmanDecoding(const std::string& encodedString, Node* root) {
    std::string decoded;
    Node* current = root;
    for (char bit : encodedString) {
        if (bit == '0') {
            current = current->left;
        }
        else { // bit == '1'
            current = current->right;
        }

        // 检查是否到达叶节点
        if (current->left == nullptr && current->right == nullptr) {
            decoded += current->symbol;
            current = root; // 重置到根节点,继续解码
        }
    }
    return decoded;
}

// 游程译码
std::string runLengthDecoding(const std::string& input) {
    std::string output;
    for (int i = 0; i < input.length(); i += 2) {
        int count = input[i] - '0';
        char symbol = input[i + 1];
        output.append(count, symbol);
    }
    return output;
}

// 模拟BSC信道
std::string simulate_BSC(const std::string& code, double error_probability) {
    std::string received = code;
    std::random_device rd;
    std::mt19937 gen(rd());
    std::bernoulli_distribution distribution(error_probability);

    for (size_t i = 0; i < received.length(); ++i) {
        if (distribution(gen)) {
            char t = '0';
            if(received[i] == t)
                received[i] = '1';
            else
                received[i] = '0';
        }
    }
    return received;
}

std::vector<int> strToBinary(const std::string& str) {
    std::vector<int> binary;
    for (char c : str) {
        binary.push_back(c - '0');
    }
    return binary;
}

std::string binaryToStr(const std::vector<int>& binary) {
    std::string str;
    for (int b : binary) {
        str += (b + '0');
    }
    return str;
}

std::vector<int> encode(const std::vector<int>& data) {
    std::vector<int> codeword(7, 0);
    for (int i = 0; i < 4; ++i) {
        codeword[i] = data[i];
    }

    // g(x) = x^3 + x + 1
    codeword[4] = (data[0] ^ data[1]) % 2;
    codeword[5] = data[1]; 
    codeword[6] = (data[0] ^ data[1] ^ data[2] ^ data[3]) % 2;

    return codeword;
}

std::vector<int> decode(const std::vector<int>& codeword) {
    return std::vector<int>(codeword.begin(), codeword.begin() + 4);
}

std::string encodeFullString(const std::string& input) {
    std::vector<int> fullEncodedData;
    for (size_t i = 0; i < input.length(); i += 4) {
        std::string substr = input.substr(i, 4);
        while (substr.length() < 4) {
            substr += '0'; 
        }
        std::vector<int> encodedData = encode(strToBinary(substr));
        fullEncodedData.insert(fullEncodedData.end(), encodedData.begin(), encodedData.end());
    }
    return binaryToStr(fullEncodedData);
}

std::string decodeFullString(const std::string& data) {
    std::vector<int> encodedData = strToBinary(data);
    std::string fullDecodedData;
    for (size_t i = 0; i < encodedData.size(); i += 7) {
        std::vector<int> codeword(encodedData.begin() + i, encodedData.begin() + i + 7);
        std::vector<int> decodedData = decode(codeword);
        fullDecodedData += binaryToStr(decodedData);
    }
    return fullDecodedData;
}

int main() {
    // 读取输入文件
    std::ifstream inputFile("in4.txt");
    std::ofstream outputFile("out4.txt");
    if (!inputFile) {
        std::cout << "Failed to open input file." << std::endl;
        return 1;
    }
    if (!outputFile) {
        std::cout << "Failed to create output file." << std::endl;
        return 1;
    }
    std::string line;
    while (std::getline(inputFile, line)) {
        if (line.empty()) {
            continue;
        }
        // 1. 游程编码结果--正确
        std::string runLengthEncoded = runLengthEncoding(line);
        outputFile << "游程编码结果:" << runLengthEncoded << std::endl;

        // 2. Huffman 编码结果--正确
        std::string huffmanEncoded = huffmanEncoding(runLengthEncoded);
        outputFile << "Huffman 编码结果:" << huffmanEncoded << std::endl;

        // 3. 加密编码结果--正确
        std::string key = generateRandomKey(huffmanEncoded.length());
        std::string encrypted = xorEncryptDecrypt(huffmanEncoded , key);
        outputFile << "加密编码结果:" << encrypted << std::endl;

        // 不是4的倍数要补0
        int len = encrypted.length();// 记录加了几个0
        while (encrypted.length() % 4 != 0) {
            encrypted += '0'; // Padding to ensure groups of 4
        }

        // 4. 信道编码结果
        std::string channelEncoded = encodeFullString(encrypted);
        outputFile << "信道编码结果:" << channelEncoded << std::endl;

        // 模拟BSC信道
        double error_probability = 0.01; // p = 0.01
        std::string received = simulate_BSC(channelEncoded, error_probability);
        outputFile << "信道传输结果: " << received << std::endl;

        // 5. 信道译码结果
        std::string channelDecoded = decodeFullString(received);
        // 要减去0
        channelDecoded = channelDecoded.substr(0, len);
        outputFile << "信道译码结果:" << channelDecoded << std::endl;

        // 6. 解密编码结果--正确
        std::string decrypted = xorEncryptDecrypt(channelDecoded, key);
        outputFile << "解密编码结果:" << decrypted << std::endl;

        // 7. Huffman 译码结果--正确
        std::string huffmanDecoded = huffmanDecoding(decrypted, pq.top());
        outputFile << "Huffman 译码结果:" << huffmanDecoded << std::endl;

        // 8. 游程译码结果--正确
        std::string runLengthDecoded = runLengthDecoding(huffmanDecoded);
        outputFile << "游程译码结果:" << runLengthDecoded << std::endl;

        // 检查解码结果是否和原文一致
        if (runLengthDecoded != line) {
            std::cout << "Error decoding line: " << line << std::endl;
        }
    }

    inputFile.close();
    outputFile.close();

    return 0;
}