CipherStreamJava/src/main/java/com/mattrixwv/cipherstream/polysubstitution/Hill.java

package com.mattrixwv.cipherstream.polysubstitution;


import java.util.ArrayList;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.mattrixwv.cipherstream.exceptions.InvalidCharacterException;
import com.mattrixwv.cipherstream.exceptions.InvalidInputException;
import com.mattrixwv.cipherstream.exceptions.InvalidKeyException;
import com.mattrixwv.matrix.ModMatrix;
import com.mattrixwv.matrix.exceptions.InvalidGeometryException;
import com.mattrixwv.matrix.exceptions.InvalidScalarException;


/**
 * Implements the Hill cipher for encoding and decoding messages using matrix operations.
 * <p>
 * The Hill cipher is a polygraphic substitution cipher that uses linear algebra to encrypt and decrypt
 * messages. It operates on blocks of text by representing them as vectors and applying matrix transformations.
 * </p>
 */
public final class Hill{
	private static final Logger logger = LoggerFactory.getLogger(Hill.class);
	//?Fields
	/** The message that needs to be encoded/decoded */
	protected String inputString;
	/** The encoded/decoded message */
	protected String outputString;
	/** The keyword used to create the grid */
	protected char characterToAdd;
	/** The matrix used perform the encoding/decoding */
	protected ModMatrix key;
	//?Settings
	/** Persist capitals in the output string */
	protected boolean preserveCapitals;
	/** Persist whitespace in the output string */
	protected boolean preserveWhitespace;
	/** Persist symbols in the output string */
	protected boolean preserveSymbols;


	/**
	 * Sets the key matrix for the Hill cipher after validating it.
	 *
	 * @param key the matrix to be used for encoding/decoding
	 * @throws InvalidKeyException if the key is not a square matrix, or does not have an inverse modulo 26
	 */
	protected void setKey(ModMatrix key) throws InvalidKeyException{
		logger.debug("Setting key");

		//Make sure the mod is correct
		logger.debug("Testing mod");
		if(key.getMod() != 26){
			throw new InvalidKeyException("This algorithm uses the english alphabet, so the mod for the key must be 26");
		}

		//Make sure the matrix is square
		logger.debug("Testing square");
		if(!key.isSquare()){
			throw new InvalidKeyException("The key must be a square matrix");
		}

		//Make sure the matrix is invertable
		logger.debug("Testing invertable");
		try{
			key.inverse();
		}
		catch(InvalidGeometryException | InvalidScalarException error){
			throw new InvalidKeyException("The key does not have an inverse mod 26");
		}

		//Set the key
		logger.debug("key\n{}", key);
		this.key = new ModMatrix(key);
	}
	/**
	 * Prepares and validates the input string for encoding by removing unwanted characters
	 * and ensuring the length is appropriate.
	 *
	 * @param inputString the message to be encoded
	 * @throws InvalidInputException if the input is null, blank, or its length is not a multiple of the matrix size
	 */
	protected void setInputStringEncode(String inputString) throws InvalidInputException{
		logger.debug("Setting input string for encoding");

		if(inputString == null){
			throw new InvalidInputException("Input must not be null");
		}

		logger.debug("Original input string '{}'", inputString);

		//Remove anything that needs removed
		if(!preserveCapitals){
			logger.debug("Removing capitals");

			inputString = inputString.toUpperCase();
		}
		if(!preserveWhitespace){
			logger.debug("Removing whitespace");

			inputString = inputString.replaceAll("[\\s]", "");
		}
		if(!preserveSymbols){
			logger.debug("Removing symbols");

			inputString = inputString.replaceAll("[^a-zA-Z\\s]", "");
		}

		//Make sure the input is correct length
		logger.debug("Checking length");
		this.inputString = inputString;
		int cleanLength = getCleanInputString().length();
		int charsToAdd = (cleanLength % key.getNumRows());
		StringBuilder inputStringBuilder = new StringBuilder();
		inputStringBuilder.append(inputString);
		if(charsToAdd != 0){
			charsToAdd = key.getNumRows() - charsToAdd;
		}
		logger.debug("Adding {} characters", charsToAdd);
		for(int cnt = 0;cnt < charsToAdd;++cnt){
			inputStringBuilder.append(characterToAdd);
		}
		inputString = inputStringBuilder.toString();

		logger.debug("Cleaned input string '{}'", inputString);
		this.inputString = inputString;

		//Make sure the input isn't blank
		if(this.inputString.isBlank() || getCleanInputString().isBlank()){
			throw new InvalidInputException("Input cannot be blank");
		}
	}
	/**
	 * Prepares and validates the input string for decoding by removing unwanted characters
	 * and ensuring the length is appropriate.
	 *
	 * @param inputString the message to be decoded
	 * @throws InvalidInputException if the input is null, blank, or its length is not a multiple of the matrix size
	 */
	protected void setInputStringDecode(String inputString) throws InvalidInputException{
		logger.debug("Setting input string for decoding");

		if(inputString == null){
			throw new InvalidInputException("Input must not be null");
		}

		logger.debug("Original input string '{}'", inputString);

		//Remove anything that needs removed
		if(!preserveCapitals){
			logger.debug("Removing capitals");

			inputString = inputString.toUpperCase();
		}
		if(!preserveWhitespace){
			logger.debug("Removing whitespace");

			inputString = inputString.replaceAll("[\\s]", "");
		}
		if(!preserveSymbols){
			logger.debug("Removing symbols");

			inputString = inputString.replaceAll("[^a-zA-Z\\s]", "");
		}

		logger.debug("Cleaned input string '{}'", inputString);
		this.inputString = inputString;

		logger.debug("Checking length");
		if(getCleanInputString().isBlank() || ((getCleanInputString().length() % key.getNumRows()) != 0)){
			throw new InvalidInputException("Length of input string must be a multiple of the number of rows in the key");
		}
	}
	/**
	 * Returns a clean version of the input string with only uppercase letters.
	 *
	 * @return the cleaned input string
	 */
	protected String getCleanInputString(){
		logger.debug("Cleaning inputString");

		String cleanInputString = inputString.toUpperCase().replaceAll("[^A-Z]", "");
		logger.debug("Clean input string '{}'", cleanInputString);

		return cleanInputString;
	}
	/**
	 * Sets the character used for padding the input string.
	 *
	 * @param characterToAdd the character to be used for padding
	 * @throws InvalidCharacterException if the character is not a letter
	 */
	protected void setCharacterToAdd(char characterToAdd) throws InvalidCharacterException{
		logger.debug("Setting character to add {}", characterToAdd);

		//Make sure the character is a letter
		if(!Character.isAlphabetic(characterToAdd)){
			throw new InvalidCharacterException("Character to add must be a letter");
		}

		//Save the characterToAdd
		if(!preserveCapitals){
			logger.debug("Removing capitals");

			characterToAdd = Character.toUpperCase(characterToAdd);
		}

		logger.debug("Cleaned character {}", characterToAdd);
		this.characterToAdd = characterToAdd;
	}
	/**
	 * Restores capitalization, whitespace, and symbols to the output string based on the original input string.
	 *
	 * @return the polished output string
	 */
	protected String polishOutputString(){
		logger.debug("Polishing output string");

		//Add the extra characters back to the output and remove the added characters
		int outputCnt = 0;
		StringBuilder outputBuilder = new StringBuilder();
		for(char ch : inputString.toCharArray()){
			logger.debug("Current char {}", ch);
			if(Character.isUpperCase(ch)){
				logger.debug("Uppercase");

				outputBuilder.append(Character.toUpperCase(outputString.charAt(outputCnt++)));
			}
			else if(Character.isLowerCase(ch)){
				logger.debug("Lowercase");

				outputBuilder.append(Character.toLowerCase(outputString.charAt(outputCnt++)));
			}
			else{
				logger.debug("Symbol");

				outputBuilder.append(ch);
			}
		}

		String cleanString = outputBuilder.toString();
		logger.debug("Polished string '{}'", cleanString);
		return cleanString;
	}
	/**
	 * Converts the cleaned input string into a list of column vectors based on the key matrix size.
	 *
	 * @return a list of vectors representing the input string
	 */
	protected ArrayList<ModMatrix> getInputVectors(){
		logger.debug("Generating input vectors");

		//Get the number of columns in the key
		int numCols = key.getNumCols();

		//Get a clean inputString
		String cleanInput = getCleanInputString();

		//Break the inputString up into lengths of numCols
		ArrayList<ModMatrix> vectors = new ArrayList<>();
		for(int cnt = 0;cnt < cleanInput.length();cnt += numCols){
			String subString = cleanInput.substring(cnt, cnt + numCols);
			int[] grid = new int[numCols];
			logger.debug("Current substring '{}'", subString);

			//Subtract 65 from each character so that A=0, B=1, ...
			for(int subCnt = 0;subCnt < subString.length();++subCnt){
				grid[subCnt] = subString.charAt(subCnt) - 65;
			}

			//Create a vector from the new values
			ModMatrix vector = new ModMatrix(26);
			vector.addCol(grid);
			logger.debug("Current vector {}", vector);

			//Add the vector to the array
			vectors.add(vector);
		}

		//Return the array of vectors
		return vectors;
	}
	/**
	 * Converts a list of vectors back into a string representation.
	 *
	 * @param outputVectors the list of vectors to be converted
	 * @return the resulting string
	 */
	protected String getOutputFromVectors(ArrayList<ModMatrix> outputVectors){
		logger.debug("Turning vectors into a string");

		//Go through each element in the vector
		StringBuilder outputBuilder = new StringBuilder();
		for(ModMatrix vector : outputVectors){
			logger.debug("Current vector {}", vector);

			//Add 65 to each element and add it to the string
			for(int cnt = 0;cnt < vector.getNumRows();++cnt){
				outputBuilder.append((char)(vector.get(cnt, 0) + 65));
			}
		}

		//Return the new string
		String convertedString = outputBuilder.toString();
		logger.debug("Converted string '{}'", convertedString);
		return convertedString;
	}
	/**
	 * Encodes the input string using the provided key matrix and stores the result in outputString.
	 */
	protected void encode(){
		logger.debug("Encoding");

		//Get an array of vectors that we are going to encode
		ArrayList<ModMatrix> inputVectors = getInputVectors();

		//Multiply the key by each vector and add the result to a new vector
		logger.debug("Multiplying vectors");
		ArrayList<ModMatrix> outputVectors = new ArrayList<>();
		for(ModMatrix inputVector : inputVectors){
			logger.debug("Current input vector {}", inputVector);
			ModMatrix outputVector = key.multiply(inputVector);

			logger.debug("Multiplied vector {}", outputVector);
			outputVectors.add(outputVector);
		}

		//Take the array of results and turn them back into letters
		outputString = getOutputFromVectors(outputVectors);

		//Add the extra characters back to the output and remove the added characters
		outputString = polishOutputString();
	}
	/**
	 * Decodes the input string using the provided key matrix and stores the result in outputString.
	 */
	protected void decode(){
		logger.debug("Decoding");

		//Get the array of vectors that we are going to decode
		ArrayList<ModMatrix> inputVectors = getInputVectors();

		//Multiply the inverse of the key by each vector and add the result to a new vector
		logger.debug("Getting inverse of key");
		ModMatrix inverseKey = key.inverse();
		logger.debug("Inverse of key {}", inverseKey);
		ArrayList<ModMatrix> outputVectors = new ArrayList<>();
		for(ModMatrix inputVector : inputVectors){
			logger.debug("Current input vector {}", inputVector);
			ModMatrix outputVector = inverseKey.multiply(inputVector);

			logger.debug("Multiplied vector {}", outputVector);
			outputVectors.add(outputVector);
		}

		//Take the array of results and turn them back into letters
		outputString = getOutputFromVectors(outputVectors);

		//Add the extra characters back to the output and remove the added characters
		outputString = polishOutputString();
	}

	//?Constructors
	/**
	 * Default constructor initializing the cipher with default settings.
	 *
	 * @throws InvalidCharacterException if the default padding character is invalid
	 */
	public Hill() throws InvalidCharacterException{
		preserveCapitals = false;
		preserveWhitespace = false;
		preserveSymbols = false;
		setCharacterToAdd('x');
		reset();
	}
	/**
	 * Constructor initializing the cipher with specified settings.
	 *
	 * @param preserveCapitals if true, preserves capitals in the output string
	 * @param preserveWhitespace if true, preserves whitespace in the output string
	 * @param preserveSymbols if true, preserves symbols in the output string
	 * @throws InvalidCharacterException if the default padding character is invalid
	 */
	public Hill(boolean preserveCapitals, boolean preserveWhitespace, boolean preserveSymbols) throws InvalidCharacterException{
		this.preserveCapitals = preserveCapitals;
		this.preserveWhitespace = preserveWhitespace;
		this.preserveSymbols = preserveSymbols;
		setCharacterToAdd('x');
		reset();
	}
	/**
	 * Constructor initializing the cipher with specified settings and padding character.
	 *
	 * @param preserveCapitals if true, preserves capitals in the output string
	 * @param preserveWhitespace if true, preserves whitespace in the output string
	 * @param preserveSymbols if true, preserves symbols in the output string
	 * @param characterToAdd the character to use for padding
	 * @throws InvalidCharacterException if the padding character is invalid
	 */
	public Hill(boolean preserveCapitals, boolean preserveWhitespace, boolean preserveSymbols, char characterToAdd) throws InvalidCharacterException{
		this.preserveCapitals = preserveCapitals;
		this.preserveWhitespace = preserveWhitespace;
		this.preserveSymbols = preserveSymbols;
		setCharacterToAdd(characterToAdd);
		reset();
	}

	/**
	 * Encodes the input string using the provided key matrix and returns the encoded string.
	 *
	 * @param key the matrix to use for encoding
	 * @param inputString the message to encode
	 * @return the encoded message
	 * @throws InvalidKeyException if the key is invalid
	 * @throws InvalidInputException if the input is invalid
	 */
	public String encode(int[][] key, String inputString) throws InvalidKeyException, InvalidInputException{
		return encode(new ModMatrix(key, 26), inputString);
	}
	/**
	 * Encodes the input string using the provided key matrix and returns the encoded string.
	 *
	 * @param key the matrix to use for encoding
	 * @param inputString the message to encode
	 * @return the encoded message
	 * @throws InvalidKeyException if the key is invalid
	 * @throws InvalidInputException if the input is invalid
	 */
	public String encode(ModMatrix key, String inputString) throws InvalidKeyException, InvalidInputException{
		setKey(key);
		setInputStringEncode(inputString);
		encode();
		return outputString;
	}
	/**
	 * Decodes the input string using the provided key matrix and returns the decoded string.
	 *
	 * @param key the matrix to use for decoding
	 * @param inputString the message to decode
	 * @return the decoded message
	 * @throws InvalidKeyException if the key is invalid
	 * @throws InvalidInputException if the input is invalid
	 */
	public String decode(int[][] key, String inputString) throws InvalidKeyException, InvalidInputException{
		return decode(new ModMatrix(key, 26), inputString);
	}
	/**
	 * Decodes the input string using the provided key matrix and returns the decoded string.
	 *
	 * @param key the matrix to use for decoding
	 * @param inputString the message to decode
	 * @return the decoded message
	 * @throws InvalidKeyException if the key is invalid
	 * @throws InvalidInputException if the input is invalid
	 */
	public String decode(ModMatrix key, String inputString) throws InvalidKeyException, InvalidInputException{
		setKey(key);
		setInputStringDecode(inputString);
		decode();
		return outputString;
	}

	/**
	 * Resets the Hill cipher by clearing all variables and setting the key to an empty matrix.
	 */
	public void reset(){
		logger.debug("Resetting fields");

		inputString = "";
		outputString = "";
		key = new ModMatrix(26);
	}

	//?Getters
	/**
	 * Returns the input string.
	 *
	 * @return the input string
	 */
	public String getInputString(){
		return inputString;
	}
	/**
	 * Returns the output string.
	 *
	 * @return the output string
	 */
	public String getOutputString(){
		return outputString;
	}
	/**
	 * Returns the key matrix used for encoding or decoding.
	 *
	 * @return the key matrix
	 */
	public ModMatrix getKey(){
		return key;
	}
}