// Arup Guha
// 9/9/2026
// Playfair cipher.

public class TestPlayFair {

	public static void main(String[] args) {
		test4();
	}
	
	// Test from class.
	public static void test1() {
		Playfair pfObj = new Playfair("MOONSTONE");
		System.out.println(pfObj);
		String code = pfObj.encrypt("SALLYWENTTOTHESHORETU");
		System.out.println(code);
		String plainBack = pfObj.decrypt(code);
		System.out.println(plainBack);
	}
	
	// Both I, J in keyword
	public static void test2() {
		Playfair pfObj = new Playfair("JELLYSILK");
		System.out.println(pfObj);
		String code = pfObj.encrypt("TELLUSWHATYOUWANTTOTALKABOUTTODAY");
		System.out.println(code);
		String plainBack = pfObj.decrypt(code);
		System.out.println(plainBack);
	}
	
	// Only I in keyword
	public static void test3() {
		Playfair pfObj = new Playfair("MATCHSTICK");
		System.out.println(pfObj);
		String code = pfObj.encrypt("THENEXTLECTUREWILLBEABOUTADVGVXAGERMANCIPHER");
		System.out.println(code);
		String plainBack = pfObj.decrypt(code);
		System.out.println(plainBack);
	}
	
	// Only J in keyword
	public static void test4() {
		Playfair pfObj = new Playfair("JACKOLANTERN");
		System.out.println(pfObj);
		String code = pfObj.encrypt("ITSNOWHERENEARHALLOWEENBUTYOUCANDOHALLOWEENHORRORNIGHTSANDPUMPKINSPICELATTESALREADY");
		System.out.println(code);
		String plainBack = pfObj.decrypt(code);
		System.out.println(plainBack);
	}
}

class Playfair {

	private char[][] box;
	
	// location[i] stores WHERE in the box letter i is.
	// So location[0] stores where 'a' is located in the grid.
	private int[] location;
	
	public Playfair(String key) {
	
		// Will store which letters are in our grid.
		boolean[] used = new boolean[26];
		
		// Allocate memory for these.
		box = new char[5][5];
		location = new int[26];
		
		// Index into box.
		int i = 0;
		
		// Copy all unique letters from keyword into box.
		for (int j=0; j<key.length(); j++) {
		
			// Current letter from key.
			char let = key.charAt(j);
			
			// Used already, don't put in grid.
			if (used[let-'A']) continue;
			
			// Store relevant key info.
			box[i/5][i%5] = let;
			location[let-'A'] = i;
			
			// Can't use this letter again.
			used[let-'A'] = true;
			
			// Special case.
			if (let == 'I') {
				used[let+1-'A'] = true;
				location[let+1-'A'] = i;
			}
			else if (let == 'J') {
				used[let-1-'A'] = true;
				location[let-1-'A'] = i;
				box[i/5][i%5] = 'I';
			}
			
			// This goes to the next square.
			i++;
		}
		
		// Now copy rest of the keyword.
		used['J'-'A'] = true;
		
		// Go through all the letters.
		for (int j=0; j<used.length; j++) {
			
			// Already placed from keyword.
			if (used[j]) continue;
			
			box[i/5][i%5] = (char)('A'+j);
			location[j] = i;
			
			// Special case for J.
			if ('A'+j == 'I') {
				location[j+1] = i;
			}
			
			// Go to the next box location.
			i++;
		}
	}
	
	// Just for testing.
	public String toString() {
		
		// Prints out box as 5 by 5 grid.
		String res = "";
		for (int i=0; i<5; i++)
			res = res + new String(box[i])+"\n";
		res = res+"\n";
		
		// Storea mapping for each letter to its location 0-based row, col.
		for (int i=0; i<26; i++) {
			res = res + (char)('A'+i)+"\t"+(location[i]/5)+", "+location[i]%5+"\n";
		}
		
		return res;
	}
	
	// Returns the result of decrypting cipher.
	public String decrypt(String cipher) {
		
		// For efficient appending...
		StringBuffer res = new StringBuffer();
		
		// Just do all pairs.
		for (int i=0; i<cipher.length(); i+=2)
			res.append(encOrDecPair(cipher.charAt(i), cipher.charAt(i+1), false));
		
		// Ta da!
		return new String(res);
	}
	
	// Returns the result of encrypting plain.
	public String encrypt(String plain) {
		
		// For efficient appending...
		StringBuffer res = new StringBuffer();
		
		int i = 0;
		while (i < plain.length()) {
			
			// Double letter.
			if ((i < plain.length()-1 && plain.charAt(i) == plain.charAt(i+1) ) ||
				(i == plain.length()-1) )	{
				
				// Padding char is X unless letter is X, then Q is used.
				char pad = plain.charAt(i) != 'X' ? 'X' : 'Q';
				
				// String we add, update index into plaintext.
				String add = encOrDecPair(plain.charAt(i), pad, true);
				res.append(add);
				i++;
				
			}
			
			// Regular case.
			else {
				String add = encOrDecPair(plain.charAt(i), plain.charAt(i+1), true);
				res.append(add);
				i+=2;
			}
			
		}
		
		return new String(res);
	}
	
	// Encrypts letters a, b.
	public String encOrDecPair(char a, char b, boolean encrypt) {
		
		// What we add to the row or col for those cases depending on
		// encrypting or decrypting.
		int add = encrypt ? 1 : 4;
		
		// Find letters on grid.
		int locA = location[a-'A'];
		int locB = location[b-'A'];
		int rowA = locA/5, colA = locA%5;
		int rowB = locB/5, colB = locB%5;
		
		// Box case.
		if (rowA != rowB && colA != colB) {
			return ""+box[rowA][colB] + box[rowB][colA];
		}
		
		// Row case.
		else if (rowA == rowB) {
			return ""+box[rowA][(colA+add)%5] + box[rowB][(colB+add)%5];
		}
		
		// Col case.
		else {
			return ""+box[(rowA+add)%5][colA] + box[(rowB+add)%5][colB];
		}
	}

}