{-# LANGUAGE RecordWildCards #-}
module CaesarCipher(
  -- * Keys
  CaesarKey,
  caesarKey,
  -- * Basic functions,
  allKeys,
  encipher,
  encipherString,
  invert,
  -- * Derived functions
  decipher,
  decipherString,
  bruteForceCaesar,
  ) where

import Data.Char
import Control.Applicative
import Math.NumberTheory.Moduli (invertMod) -- package arithmoi
import Data.Maybe (catMaybes)

----------- Helpers

textToDigits::String->[Integer]
textToDigits = map (\x->toInteger (ord x - 97)) 

digitsToText::[Integer]->String
digitsToText = map (\x->chr (fromIntegral x + 97)) 

----------- Constructing Keys

-- invariant: r and q are inverses mod 26. 
-- To ensure this invariant, we only export the 'caesarKey' smart constructor
-- and not the 'Key' constructor
data CaesarKey = Key { r :: Integer, s :: Integer, q :: Integer }

caesarKey :: Integer -> Integer -> Maybe CaesarKey
caesarKey r s = Key r s <$> invertMod r 26

------------ Basic functions
allKeys :: [CaesarKey]
allKeys = catMaybes $ caesarKey <$> [0..25] <*> [0..25]

encipher :: CaesarKey -> Integer -> Integer
encipher Key{..} p = mod (r * p + s) 26

encipherString :: CaesarKey -> String -> String
encipherString key = digitsToText . map (encipher key) . textToDigits
    
-- turns a key suitable for encoding into one suitable for decoding
invert :: CaesarKey -> CaesarKey
invert (Key r s q) = Key q ((26-q)*s) r

------------ Optional "derived" functions
decipher :: CaesarKey -> Integer -> Integer
decipher key = encipher (invert key)

decipherString :: CaesarKey -> String -> String
decipherString key = encipherString (invert key)

bruteForceCaesar :: String -> [String]
bruteForceCaesar str = [decipherString key str | key <- allKeys]