#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <algorithm>
#include <cctype>

// ADFGVX字母表
const std::string ADFGVX_CHARS = "ADFGVX";

// 固定的波利比奥斯方阵 (6x6)
// 包含26个字母和10个数字
const char POLYBIUS_SQUARE[6][6] = {
    {'P', 'H', 'R', 'A', 'S', 'E'},
    {'B', 'C', 'D', 'F', 'G', 'I'},
    {'J', 'K', 'L', 'M', 'N', 'O'},
    {'Q', 'T', 'U', 'V', 'W', 'X'},
    {'Y', 'Z', '0', '1', '2', '3'},
    {'4', '5', '6', '7', '8', '9'}
};

// 用于快速查找字符坐标的映射
std::map<char, std::pair<int, int>> charToCoords;

// 初始化字符到坐标的映射
void initializeCharToCoords() {
    for (int i = 0; i < 6; ++i) {
        for (int j = 0; j < 6; ++j) {
            charToCoords[POLYBIUS_SQUARE[i][j]] = {i, j};
        }
    }
}

// 将明文替换为ADFGVX坐标对
std::string substitute(const std::string& plaintext) {
    std::string substitutedText = "";
    for (char c : plaintext) {
        c = toupper(c); // 转换为大写
        if (charToCoords.count(c)) {
            substitutedText += ADFGVX_CHARS[charToCoords[c].first];
            substitutedText += ADFGVX_CHARS[charToCoords[c].second];
        } else if (isdigit(c)) { // 处理数字
            // 查找数字在方阵中的位置
            bool found = false;
            for (int i = 0; i < 6; ++i) {
                for (int j = 0; j < 6; ++j) {
                    if (POLYBIUS_SQUARE[i][j] == c) {
                        substitutedText += ADFGVX_CHARS[i];
                        substitutedText += ADFGVX_CHARS[j];
                        found = true;
                        break;
                    }
                }
                if (found) break;
            }
        } else {
            // 忽略非字母数字字符
        }
    }
    return substitutedText;
}

// 根据换位密钥进行列换位加密
std::string transposeEncrypt(const std::string& substitutedText, const std::string& key) {
    std::string cleanKey = "";
    for (char c : key) {
        if (isalpha(c)) {
            cleanKey += toupper(c);
        }
    }

    int keyLength = cleanKey.length();
    if (keyLength == 0) return substitutedText; // No transposition if key is empty

    int numRows = (substitutedText.length() + keyLength - 1) / keyLength;
    std::vector<std::vector<char>> grid(numRows, std::vector<char>(keyLength, ' '));

    int textIndex = 0;
    for (int r = 0; r < numRows; ++r) {
        for (int c = 0; c < keyLength; ++c) {
            if (textIndex < substitutedText.length()) {
                grid[r][c] = substitutedText[textIndex++];
            }
        }
    }

    // 创建密钥索引对，用于排序
    std::vector<std::pair<char, int>> keyOrder(keyLength);
    for (int i = 0; i < keyLength; ++i) {
        keyOrder[i] = {cleanKey[i], i};
    }
    std::sort(keyOrder.begin(), keyOrder.end());

    std::string ciphertext = "";
    for (int i = 0; i < keyLength; ++i) {
        int originalCol = keyOrder[i].second;
        for (int r = 0; r < numRows; ++r) {
            if (grid[r][originalCol] != ' ') {
                ciphertext += grid[r][originalCol];
            }
        }
    }
    return ciphertext;
}

// 根据换位密钥进行列逆换位解密
std::string transposeDecrypt(const std::string& ciphertext, const std::string& key) {
    std::string cleanKey = "";
    for (char c : key) {
        if (isalpha(c)) {
            cleanKey += toupper(c);
        }
    }

    int keyLength = cleanKey.length();
    if (keyLength == 0) return ciphertext; // No transposition if key is empty

    int numRows = (ciphertext.length() + keyLength - 1) / keyLength;
    std::vector<std::vector<char>> grid(numRows, std::vector<char>(keyLength, ' '));

    // 计算每列的长度
    std::vector<int> colLengths(keyLength);
    int remaining = ciphertext.length();
    for (int i = 0; i < keyLength; ++i) {
        colLengths[i] = numRows;
        if (remaining % keyLength != 0 && i >= (keyLength - (numRows * keyLength - ciphertext.length()))) {
            colLengths[i]--; // 最后一行的空位
        }
        remaining -= colLengths[i];
    }

    // 重新构建密钥索引对，用于排序
    std::vector<std::pair<char, int>> keyOrder(keyLength);
    for (int i = 0; i < keyLength; ++i) {
        keyOrder[i] = {cleanKey[i], i};
    }
    std::sort(keyOrder.begin(), keyOrder.end());

    // 填充网格
    int currentCipherIndex = 0;
    for (int i = 0; i < keyLength; ++i) {
        int originalCol = keyOrder[i].second;
        for (int r = 0; r < colLengths[originalCol]; ++r) {
            grid[r][originalCol] = ciphertext[currentCipherIndex++];
        }
    }

    std::string substitutedText = "";
    for (int r = 0; r < numRows; ++r) {
        for (int c = 0; c < keyLength; ++c) {
            if (grid[r][c] != ' ') {
                substitutedText += grid[r][c];
            }
        }
    }
    return substitutedText;
}

// 将ADFGVX坐标对逆替换回明文
std::string substituteDecrypt(const std::string& substitutedText) {
    std::string plaintext = "";
    for (size_t i = 0; i < substitutedText.length(); i += 2) {
        char rowChar = substitutedText[i];
        char colChar = substitutedText[i+1];

        int row = ADFGVX_CHARS.find(rowChar);
        int col = ADFGVX_CHARS.find(colChar);

        if (row != std::string::npos && col != std::string::npos) {
            plaintext += POLYBIUS_SQUARE[row][col];
        }
    }
    return plaintext;
}

int main() {
    initializeCharToCoords(); // 初始化映射

    std::string plaintext;
    std::string transpositionKey;

    std::cout << "=====================================" << std::endl;
    std::cout << "        ADFGVX密码加密/解密程序" << std::endl;
    std::cout << "=====================================" << std::endl;

    std::cout << "\n请输入要加密的消息 (只包含字母和数字): ";
    std::getline(std::cin, plaintext);

    std::cout << "请输入换位密钥 (只包含字母): ";
    std::getline(std::cin, transpositionKey);

    // 加密过程
    std::string substituted = substitute(plaintext);
    std::string encrypted = transposeEncrypt(substituted, transpositionKey);
    std::cout << "\n加密后的密文是: " << encrypted << std::endl;

    // 解密过程
    std::string decryptedSubstituted = transposeDecrypt(encrypted, transpositionKey);
    std::string decryptedPlaintext = substituteDecrypt(decryptedSubstituted);
    std::cout << "解密后的明文是: " << decryptedPlaintext << std::endl;

    return 0;
}