Keyed caesar

open Printf

module Alphabets = struct
  (*let original_latin_alphabet = "ABCDEFZHIKLMNOPQRSTVX"*)
  (*let old_latin_alphabet = "ABCDEFGHIKLMNOPQRSTVX"*)
  let classical_latin_alphabet = "ABCDEFGHIKLMNOPQRSTVXYZ"
  (*let modern_latin_alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"*)
end

let deduplicate str =
  let result = Buffer.create (String.length str) in
  let tbl = Hashtbl.create (String.length str) in
  String.iter (
    fun c ->
      if not (Hashtbl.mem tbl c) then begin
        Buffer.add_char result c;
        Hashtbl.add tbl c ();
      end
  ) str;
  Buffer.contents result

let filter_string s f =
  let buf = Buffer.create (String.length s) in
  String.iter (fun c -> if f c then Buffer.add_char buf c) s;
  Buffer.contents buf

let keyed_alphabet keyword alphabet =
  let dedup_keyword = deduplicate (String.uppercase_ascii keyword) in
  let alpha_minus_keyword =
    filter_string alphabet (fun c -> not (String.contains dedup_keyword c))
  in
  dedup_keyword ^ alpha_minus_keyword

let caesar_encode keyed_alphabet shift text =
  let shift_mod = shift mod String.length keyed_alphabet in
  let find_index c =
    try Some (String.index keyed_alphabet c)
    with Not_found -> None
  in
  let shift_char c =
    match find_index c with
    | Some index ->
        keyed_alphabet.[(index + shift_mod) mod String.length keyed_alphabet]
    | None -> c
  in
  String.map shift_char text

let caesar_decode keyed_alphabet shift text =
  let shift_mod = shift mod String.length keyed_alphabet in
  let find_index c =
    try Some (String.index keyed_alphabet c)
    with Not_found -> None
  in
  let shift_char c =
    match find_index c with
    | Some index ->
        let shifted_index = (index - shift_mod) mod String.length keyed_alphabet in
        let corrected_index =
          if shifted_index < 0
            then shifted_index + String.length keyed_alphabet
          else
            shifted_index in
              keyed_alphabet.[corrected_index]
    | None -> c
  in
  String.map shift_char text

let () =
  let keyword = "helo" in
  let text_to_encrypt = "HELLO, WORLD!" in
  let shift = 13 in (* Shift by 13, as in the standard Caesar cipher *)
  let keyed_alpha = keyed_alphabet keyword Alphabets.classical_latin_alphabet in
  let encrypted_text = caesar_encode keyed_alpha shift text_to_encrypt in
  let decrypted_text = caesar_decode keyed_alpha shift encrypted_text in

  printf "Original text: %s\n" text_to_encrypt;
  printf "Keyword: %s\n" keyword;
  printf "Keyed alphabet: %s\n" keyed_alpha;
  printf "Shifted by: %d\n" shift;
  printf "Encrypted text: %s\n" encrypted_text;
  printf "Decrypted text: %s\n" decrypted_text