module Subst = struct
  module Table = Map.Make (Char)

  let make_table cs = Table.of_seq cs

  let is_alpha c =
    let open Char in
    let cc = code (uppercase_ascii c) in
    cc >= code 'A' && cc <= code 'Z'

  let subst t s =
    String.map
      (fun c ->
        let table = make_table t in
        if is_alpha c then
          Option.value
            (Table.find_opt (Char.uppercase_ascii c) table)
            ~default:'_'
        else c)
      s
end

module Caesar = struct
  let alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
  let i c = Char.code c - Char.code 'A'

  let rot_table n =
    let open Char in
    Seq.map
      (fun c -> (c, chr (((i c + n) mod 26) + code 'A')))
      (String.to_seq alphabet)
  
  let enc n = Subst.subst (rot_table n)
  let dec n = Subst.subst (rot_table (26 - n))

end

module Ex1 = struct
  open Caesar

  let brute s =
    for n = 1 to 26 do
      Printf.printf "%2d: %s\n" n (dec n s)
    done
end

module Ex2 = struct
  let table =
    [
      ('A', 'N');
      ('B', 'D');
      ('C', 'O');
      ('D', 'P');
      ('E', 'E');
      ('F', 'Q');
      ('G', 'R');
      ('H', 'S');
      ('I', 'T');
      ('J', 'U');
      ('K', 'B');
      ('L', 'V');
      ('M', 'W');
      ('O', 'A');
      ('Q', 'Z');
      ('R', 'F');
      ('S', 'G');
      ('T', 'C');
      ('U', 'H');
      ('V', 'I');
      ('W', 'J');
      ('X', 'K');
      ('Y', 'L');
      ('Z', 'M');
    ]
  let dec = Subst.subst (List.to_seq table) 
end

module Ex3 = struct
  let table =
    [
      ('A', 'W');
      ('B', 'A');
      ('C', 'G');
      ('D', 'X');
      ('E', 'B');
      ('F', 'Y');
      ('G', 'Z');
      ('H', 'H');
      ('I', 'I');
      ('J', 'J');
      ('K', 'K');
      ('L', 'F');
      ('M', 'L');
      ('N', 'M');
      ('O', 'E');
      ('P', 'N');
      ('Q', 'O');
      ('R', 'C');
      ('S', 'P');
      ('T', 'D');
      ('U', 'Q');
      ('V', 'R');
      ('W', 'S');
      ('X', 'T');
      ('Y', 'U');
      ('Z', 'V');
    ]
  let dec = Subst.subst (List.to_seq table) 
end