(ns cipher-analytical-machine.ciphers.gamma
  (:require
   [clojure.set :as set]
   [clojure.string :as cs]
   [cipher-analytical-machine.ciphers.simple-substitution :as ss])
  (:gen-class))

(defn add-mod
  [f s module]
  (-> (+ f s)
      (mod module)))

(defn generate-seq
  "Generate a lazy sequence for a key (a b c)."
  [a b c module]
  ((fn build-seq [a b c]
     (lazy-seq (cons a (build-seq b c (add-mod a c module))))) a b c))

(defn generate-gamma-seq
  "Generate the gamma seq from a sequence."
  [acc seq module]
  (if (= (second seq) nil) acc
    (generate-gamma-seq
     (conj acc (add-mod (first seq) (second seq) module))
     (rest seq)
     module)))

(defn generate-gamma
  "Generate the gamma from a key (a b c)."
  [a b c module size]
  (generate-gamma-seq []
                      (->> (generate-seq a b c module)
                           (take (inc size)))
                      module))

(defn encrypt-array
  "Encrypt an array using the gamma cipher. The function is case-insensitive. The key is an array of three elements [a b c]. The array contains integer in range [0, module)."
  [key module array]
  (let [size (count array)
        gamma (->> (apply conj [key module size])
                   (apply generate-gamma))]
    (->> gamma
         (map vector array)
         (map (fn [[a g]] (add-mod a g module))))))

(defn encrypt-message
  "Encrypt a message using the gamma cipher. The function is case-insensitive. The key is an array of three elements [a b c]. The array contains integer in range [0, module)."
  [key symbols message]
  (let [module (count symbols)
        table (ss/generate-sorted-substitution-table symbols)]
    (->> message
         (ss/decrypt-by-table (set/map-invert table))
         (encrypt-array key module)
         (ss/decrypt-by-table table)
         (cs/join))))

(defn decrypt-array
  "Decrypt an array using the gamma cipher. The function is case-insensitive. The key is an array of three elements [a b c]. The array contains integer in range [0, module)."
  [key module array]
  (let [size (count array)
        gamma (->> (apply conj [key module size])
                   (apply generate-gamma))]
    (->> gamma
         (map -)
         (map vector array)
         (map (fn [[a g]] (add-mod a g module))))))