import java.util.Scanner;

public class PlayfairCipher {

    static char[][] matrix = new char[5][5];

    static void createMatrix(String key) {
        key = key.toUpperCase().replaceAll("[^A-Z]", "").replace("J", "I");

        boolean[] used = new boolean[26];
        used['J' - 'A'] = true;

        StringBuilder chars = new StringBuilder();

        for (char ch : key.toCharArray()) {
            if (!used[ch - 'A']) {
                chars.append(ch);
                used[ch - 'A'] = true;
            }
        }

        for (char ch = 'A'; ch <= 'Z'; ch++) {
            if (!used[ch - 'A']) {
                chars.append(ch);
                used[ch - 'A'] = true;
            }
        }

        int index = 0;

        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < 5; j++) {
                matrix[i][j] = chars.charAt(index++);
            }
        }
    }

    static String prepareText(String text) {
        text = text.toUpperCase().replaceAll("[^A-Z]", "").replace("J", "I");

        StringBuilder result = new StringBuilder();

        int i = 0;

        while (i < text.length()) {

            char first = text.charAt(i);

            if (i + 1 < text.length()) {
                char second = text.charAt(i + 1);

                if (first == second) {
                    result.append(first);
                    result.append('X');
                    i++;
                } else {
                    result.append(first);
                    result.append(second);
                    i += 2;
                }
            } else {
                result.append(first);
                result.append('X');
                i++;
            }
        }

        return result.toString();
    }

    // Find row and column of a character
    static int[] findPosition(char ch) {

        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < 5; j++) {

                if (matrix[i][j] == ch) {
                    return new int[]{i, j};
                }
            }
        }

        return null;
    }

    // Encrypt plaintext
    static String encrypt(String text) {

        StringBuilder encrypted = new StringBuilder();

        for (int i = 0; i < text.length(); i += 2) {

            char first = text.charAt(i);
            char second = text.charAt(i + 1);

            int[] pos1 = findPosition(first);
            int[] pos2 = findPosition(second);

            int row1 = pos1[0];
            int col1 = pos1[1];

            int row2 = pos2[0];
            int col2 = pos2[1];

            // Same row
            if (row1 == row2) {

                encrypted.append(matrix[row1][(col1 + 1) % 5]);
                encrypted.append(matrix[row2][(col2 + 1) % 5]);

            }

            // Same column
            else if (col1 == col2) {

                encrypted.append(matrix[(row1 + 1) % 5][col1]);
                encrypted.append(matrix[(row2 + 1) % 5][col2]);

            }

            else {

                encrypted.append(matrix[row1][col2]);
                encrypted.append(matrix[row2][col1]);
            }
        }

        return encrypted.toString();
    }

    static void displayMatrix() {

        System.out.println("\nPlayfair Matrix:");

        for (int i = 0; i < 5; i++) {

            for (int j = 0; j < 5; j++) {
                System.out.print(matrix[i][j] + " ");
            }

            System.out.println();
        }
    }

    public static void main(String[] args) {

        Scanner sc = new Scanner(System.in);

        System.out.print("Enter key: ");
        String key = sc.nextLine();

        System.out.print("Enter plaintext: ");
        String plaintext = sc.nextLine();

        createMatrix(key);

        displayMatrix();

        String preparedText = prepareText(plaintext);

        System.out.println("\nPrepared Plaintext: " + preparedText);

        String encryptedText = encrypt(preparedText);

        System.out.println("Encrypted Text: " + encryptedText);

        sc.close();
    }
}