{- | First assignment for IB016, semester spring 2017, 20 points (+bonus).

== Task overview

Your task is to implement several simple encryption/decryption routines.

  * Monoalphabetic substitution cipher, for general description see the
    [Wikipedia page on substitution cipher](http://en.wikipedia.org/wiki/Substitution_cipher).
  * Caesar shift cipher, for general description see the
    [Wikipedia page on Caesar cipher](http://en.wikipedia.org/wiki/Caesar_cipher).
  * Polyalphabetic substitution encryption called Vigenere cipher, for general description see
    [Wikipedia page on Vigenere cipher](http://en.wikipedia.org/wiki/Vigen%C3%A8re_cipher).

The ciphers you should implement differ slightly from their historical
prototypes. Namely, they work only on lower-case English letters and ignore
other characters, letters are numbered @0..25@. Please make sure to adhere to
the specification provided below.

== Tips and tricks

  * You may find several useful functions in
  [Prelude](https://hackage.haskell.org/package/base-4.8.2.0/docs/Prelude.html),
  [Data.List](https://hackage.haskell.org/package/base-4.8.2.0/docs/Data-List.html),
  [Data.Char](https://hackage.haskell.org/package/base-4.8.2.0/docs/Data-Char.html) and
  [Data.Tuple](https://hackage.haskell.org/package/base-4.8.2.0/docs/Data-Tuple.html).
  * The 'mod' function does not guarantee a positive result.

== Module and package constraints

You can use any modules from the
<https://hackage.haskell.org/package/base-4.8.2.0/ base> package if you wish.
However, try not to use list indexing, i.e. do not use the function '!!'.  If
you wish so, you can also use Unicode syntax from
<https://hackage.haskell.org/package/unicode-prelude-0.1.1 unicode-prelude> but
no other packages are allowed.
-}

-- Name: Jan Pokorný
-- UID: 456195

module Ciphers (
    -- * Monoaplhabetic substitution cipher (4 points)
      encryptSubstitution
    , decryptSubstitution
    -- * Caesar cipher (3 points)
    , encryptCaesar
    , decryptCaesar
    -- * Vigenere cipher (4 points)
    , encryptVigenere
    , decryptVigenere
    -- * Using multiple ciphers (3 points)
    , applyMultiple
    -- * Cryptanalysis assistant (6+ points)
    , assistant
) where
import Data.List
import Data.Maybe
import Data.Tuple
import Data.Char
import Text.Printf

-- UTILITY FUNCTIONS --
-- ASCII value of 'a'
baseChar :: Int
baseChar = fromEnum 'a'

-- alphabet length
numChars :: Int
numChars = 26

-- Char to 0-25
c2i :: Char -> Int
c2i = subtract baseChar . fromEnum

-- 0-25 (mod 26) to Char
i2c :: Int -> Char
i2c = toEnum . (+baseChar) . (`mod` numChars) . (+numChars) . (`mod` numChars)

-- shifts every character by @shift@ places, can be negative
stringShift :: Int -> String -> String
stringShift shift = map $ i2c . (+shift) . c2i

-- performs substituion without checking permutation validity
substitute :: [(Char, Char)] -> String -> String
substitute permutation = map (\c -> fromMaybe c $ lookup c permutation)

-- DATA --
-- English letter frequency
enFreq :: [(Char, Double)]
enFreq = [('a', 8.167), ('b', 1.492), ('c', 2.782), ('d', 4.253), ('e', 2.702), ('f', 2.228), ('g', 2.015), ('h', 6.094), ('i', 6.966), ('j', 0.153), ('k', 0.772), ('l', 4.025), ('m', 2.406), ('n', 6.749), ('o', 7.507), ('p', 1.929), ('q', 0.095), ('r', 5.987), ('s', 6.327), ('t', 9.056), ('u', 2.758), ('v', 0.978), ('w', 2.360), ('x', 0.150), ('y', 1.974), ('z', 0.074)]


-- | Encrypt the given plaintext using monoaplhabetic substitution according to
-- the permutation given in the for of associative list. In this case the range
-- of characters is not limited (any value of 'Char' type can map to any value
-- of 'Char' type).
--
-- The first argument is a list of 'Char' tuples representing the permutation.
-- The encryption process works on each character of the plaintext independently
-- in the following way:
--
--   * If the list contains @(x, y)@ then @x@ is encrypted as @y@.
--   * If there is no tuple with @x@ as the first component, @x@ is encrypted as @x@.
--
-- The valid permutation is defined as follows:
-- 
--  * If the list contains @(x, y)@, then it contains @(y, z)@.
--  * If the list contains @(x, z)@ and @(y, z)@, then @x == y@.
--  * If the list contains @(x, y)@ and @(x, z)@, then @y == z@.
--  * The list does not contain any element multiple times.
--  * The set of characters in the first components is the same as
--    in the second components of the tuples.
--
-- If the first argument is not a permutation the function should fail
-- with an appropriate error message.
--
-- >>> encryptSubstitution [('a', 'b'), ('b', 'c'), ('c', 'a')] "abcde"
-- "bcade"
--
-- >>> encryptSubstitution [('a', 'b'), ('b', 'c'), ('c', 'd')] "abcde"
-- *** Exception: Not a permutation.
--
-- >>> encryptSubstitution [] "abcde"
-- "abcde"
encryptSubstitution :: [(Char, Char)] -- ^ Character permutation
                    -> String         -- ^ Plaintext
                    -> String         -- ^ Ciphertext
encryptSubstitution permutation
  | validPermutation permutation = substitute permutation
  | otherwise                    = error "Not a permutation."
  where validPermutation perm = r5ValidAll perm && and (allValid <$> [r1ValidTuple, r2ValidTuple, r3ValidTuple, r4ValidTuple] <*> pure perm)
        -- Rules for valid perm, in the same order as in doc text
        allValid rule xs = all (`rule` xs) xs
        r1ValidTuple (_, y) = any ((==y) . fst)
        r2ValidTuple (x, z) = all (\(y, zz) -> (z /= zz) || (x == y))
        r3ValidTuple (x, y) = all (\(xx, z) -> (x /= xx) || (y == z))
        r4ValidTuple t = (==1) . length . filter (==t)
        r5ValidAll xs = sort (map fst xs) == sort (map snd xs)

-- | Decrypt the given plaintext using monoaplhabetic substitution.
-- For detailed specification, see 'encryptSubstitution'.
--
-- For every valid permutation @p@ and for every String @s@ it should
-- hold that @decryptSubstitution p (encryptSubstitution p s) == s@.
--
-- >>> decryptSubstitution [('a', 'b'), ('b', 'c'), ('c', 'a')] "bcade"
-- "abcde"
--
-- >>> decryptSubstitution [('a', 'b'), ('b', 'c'), ('c', 'd')] "bcade"
-- *** Exception: Not a permutation.
decryptSubstitution :: [(Char, Char)] -- ^ Character permutation
                    -> String         -- ^ Ciphertext
                    -> String         -- ^ Plaintext
decryptSubstitution perm = encryptSubstitution (map swap perm)

-- | Shifts each input character by adding given shift value (modulo 26). This
-- works only on English lowercase letters. That is, if @'a'@ is shifted by 3,
-- it is changed to @'d'@, @'z'@ shifted by 3 is @'c'@.
--
-- When shifting (modulo 26) you can make use of function
-- [fromEnum](http://hackage.haskell.org/packages/archive/base/latest/doc/html/Prelude.html#v:fromEnum)
-- which converts any 'Enum' values (including 'Char') to number. There is also
-- 'toEnum' function defined in the same
-- [typeclass](http://hackage.haskell.org/packages/archive/base/latest/doc/html/Prelude.html#t:Enum).
--
-- The shift must be positive number, but it need not be less then 26.
--
-- >>> encryptCaesar 5 "abc"
-- "fgh"
--
-- >>> encryptCaesar (-3) "abc"
-- *** Exception: Incorrect shift.
encryptCaesar :: Int    -- ^ Shift
              -> String -- ^ Plaintext
              -> String -- ^ Ciphertext
encryptCaesar shift
  | shift >= 0 = stringShift shift
  | otherwise  = error "Incorrect shift."

-- | Decrypt the given plaintext using Caesar cipher.
-- For detailed specification, see 'encryptCaesar'.
--
-- For every valid shift 'i' and for every plaintext string 's' it should hold
-- that @decryptCaesar i (encryptCaesar i s) == s@.
--
-- >>> decryptCaesar 5 "fgh"
-- "abc"
--
-- >>> decryptCaesar (-3) "abc"
-- *** Exception: Incorrect shift.
decryptCaesar :: Int    -- ^ Shift
              -> String -- ^ Plaintext
              -> String -- ^ Ciphertext
decryptCaesar shift
  | shift >= 0 = stringShift (-shift)
  | otherwise  = error "Incorrect shift."

-- | Adds characters from the input 'String' and from key 'String' modulo 26.
-- It works only on characters @'a'..'z'@.  If the key string is shorter then
-- input string it is repeated.
--
-- Similarly to 'encryptCaesar' you can use 'fromEnum' and 'toEnum' to calculate character codes. For example, sum of @'b'@ and @'c'@ gives @'d'@.
--
-- As a special case, if the key is empty the result is the same as input.
--
-- >>>  encryptVigenere "gfe" "abcdefg"
-- "gggjjjm"
--
-- >>> encryptVigenere "heslo" "hello world"
-- "oidwc agczk"
--
-- >>> encryptVigenere "" "hello world"
-- "hello world"
encryptVigenere :: String -- ^ Key string
                -> String -- ^ Plaintext
                -> String -- ^ Ciphertext
encryptVigenere "" = id
encryptVigenere key = zipWith encChar (cycle key)
  where encChar k p
          | p `elem` ['a'..'z'] = i2c (c2i k + c2i p)
          | otherwise           = p

-- | Decrypt the given plaintext using Vigenere cipher.
-- For detailed specification, see 'encryptVigenere'.
--
-- For every valid key string @k@ and for every plaintext string @s@ it should
-- hold that @decryptVigenere k (encryptVigenere k s) == s@.
--
-- >>> decryptVigenere "abcd" "abcd"
-- "aaaa"
--
-- >>> decryptVigenere "" "hello"
-- "hello"
decryptVigenere :: String -- ^ Key string
                -> String -- ^ Plaintext
                -> String -- ^ Ciphertext
decryptVigenere key = encryptVigenere (map (i2c . (numChars-) . c2i) key)

-- | Combines several ciphers into one, using each one for a block of given
-- size. For example @applyMultiple [f, g] 4 plain@ applies cipher @f@ to first 4
-- characters of @plain@, cipher @g@ to next 4, then again @f@… The last block
-- can be shorter than given number (length of plaintext need not be multiple
-- of block length).
--
-- The first argument is list of encryption (or decryption) functions, these
-- function are used for encryption (or decryption) of each blocks of the input
-- text (each block has length given by the second argument). Functions are
-- repeated in the same way as the key for Vigenere cipher.
--
-- * In the case the length of the input string is not divisible by block
-- length the last block can be shorter.
--
-- * In the case the block length is less or equal to zero, the resulting
-- string is same as the input string.
--
-- * In the case the list of functions is empty, the resulting string is the
-- same as the input string.
--
-- >>> applyMultiple [encryptCaesar 1, encryptCaesar 2, encryptCaesar  3] 2 "aaaaaaaaaaaaa"
-- "bbccddbbccddb"
--
-- >>> applyMultiple [encryptVigenere "abc", encryptVigenere "bce", encryptCaesar 5] 4 "abcdefghijklmnopq"
-- "acedfhkinopqmoqpr"
--
-- >>> applyMultiple [decryptVigenere "abc", decryptVigenere "bce", decryptCaesar 5] 4 "acedfhkinopqmoqpr"
-- "abcdefghijklmnopq"
--
-- >>> applyMultiple [] 4 "hello"
-- "hello"
--
-- >>> applyMultiple [encryptCaesar 2, encryptCaesar 12] (-3) "hello"
-- "hello"
applyMultiple :: [String -> String] -- ^ List of encryption functions
              -> Int                -- ^ Block size
              -> String             -- ^ Plaintext
              -> String             -- ^ Ciphertext
applyMultiple [] _ plain = plain
applyMultiple fns blockSize plain
  | blockSize > 0 = concat $ zipWith ($) (cycle fns) (blockify blockSize plain)
  | otherwise = plain
  where blockify _ "" = []
        blockify size str = take size str : blockify size (drop size str)

-- | Interactive cryptanalysis assistant for monoalphabetic substitution cipher.
--
-- Minimal requirements:
--
--   * The assistant displays current state of plaintext/ciphertext
--   * The assistant displays currently applied substitutions
--   * You can add a new substitution
--   * You can delete a substitution
--   * You can delete all substitutions at once
--   * You can close the assistant returning the final substitution
--   * The substitution currently applied need not be correct permuation as
--   defined in 'encryptSubstitution'
--
-- The outputs below are just an example, you don't have to adhere
-- to that syntax (as long as you support the minimal functionality above).
--
-- @
-- >>>  assistant "abcdef"
--
-- Current state of decryption (CIPHERTEXT, plaintext):
-- ABCDEF
-- Current substitutions:
--
-- Available commands:
--   a (add new substitution pair)
--   d (delete substitution pair)
--   r (reset current substitution)
--   q (quit assistant)
-- Command choice: a
-- Ciphertext letter: a
-- Plaintext letter: h
--
-- Current state of decryption (CIPHERTEXT, plaintext):
-- hBCDEF
-- Current substitutions: h->A
--
-- Available commands:
--   a (add new substitution pair)
--   d (delete substitution pair)
--   r (reset current substitution)
--   q (quit assistant)
-- Command choice: a
-- Ciphertext letter: b
-- Plaintext letter: e
--
-- Current state of decryption (CIPHERTEXT, plaintext):
-- heCDEF
-- Current substitutions: e->B, h->A
--
-- Available commands:
--   a (add new substitution pair)
--   d (delete substitution pair)
--   r (reset current substitution)
--   q (quit assistant)
-- Command choice: r
--
-- Current state of decryption (CIPHERTEXT, plaintext):
-- ABCDEF
-- Current substitutions:
--
-- Available commands:
--   a (add new substitution pair)
--   d (delete substitution pair)
--   r (reset current substitution)
--   q (quit assistant)
-- Command choice: q
-- []
-- @
--
-- == BONUS
--
-- Add any other functions to the assistant you find useful.
-- You can get up to 5 additional point depending on the functionality
-- you implement.
assistant :: String            -- ^ Ciphertext
          -> IO [(Char, Char)] -- ^ Final substitution
assistant str = assist [] $ map toUpper str
  where assist :: [(Char, Char)] -> String -> IO [(Char, Char)]
        assist substitutions ciphertext = do
          putStrLn "Current state of decryption (CIPHERTEXT, plaintext):"
          putStrLn $ substitute substitutions ciphertext
          putStrLn "Current substitutions: "
          putStrLn $ showSubstitutions substitutions
          putStr "\n\
                  \Available commands:\n\
                  \  a (add new substitution pair)\n\
                  \  d (delete substitution pair)\n\
                  \  r (reset current substitution)\n\
                  \  f (frequency analysis)\n\
                  \  e (show reference English frequencies)\n\
                  \  c (show recommendations based on frequency in English)\n\
                  \  q (quit assistant)\n\
                  \Command choice: "
          command <- getLine
          failEmpty command
          case head command of
            'a' -> do
                    putStrLn "Adding new pair..."
                    (c, p) <- selectPair
                    let colliding = filter (\(cc, pp) -> cc == c || pp == p) substitutions
                    if null colliding
                      then assist (sort ((c, p):substitutions)) ciphertext
                      else do
                        putStrLn $ "Can't assign, colliding with: \n" ++ showSubstitutions colliding
                        assist substitutions ciphertext
            'd' -> do
                    putStrLn "Deleting pair..."
                    pair <- selectPair
                    assist (delete pair substitutions) ciphertext
            'r' -> do
                    putStrLn "Resetting..."
                    assist [] ciphertext
            'f' -> do
                    putStrLn "Doing frequency analysis..."
                    putStrLn $ showFrequencies $ getFrequency ciphertext
                    assist substitutions ciphertext
            'e' -> do 
                    putStrLn "English letter frequency:"
                    putStrLn $ showFrequencies enFreq
                    assist substitutions ciphertext
            'c' -> do
                    putStrLn "Recommending pairings based on English frequency..."
                    let cipherFreq = getFrequency ciphertext
                    let (assigned_c, assigned_p) = unzip substitutions
                    let (unassigned_c, unassigned_p) = (['a'..'z'] \\ assigned_c, ['a'..'z'] \\ assigned_p)
                    let getFreq freq c = fromMaybe 0.0 $ lookup c freq
                    let compareFreq freq c d = compare (getFreq freq c) (getFreq freq d)
                    let proposedPairs = zip (sortBy (compareFreq cipherFreq) unassigned_c) (sortBy (compareFreq enFreq) unassigned_p)
                    putStrLn $ showSubstitutions proposedPairs
                    assist substitutions ciphertext

            'q' -> return substitutions
            _   -> assist substitutions ciphertext
        
        selectPair = do
          putStr "Ciphertext letter: "
          c <- getLine
          failEmpty c
          putStr "Plaintext letter: "
          p <- getLine
          failEmpty p
          return (toUpper $ head c, toLower $ head p)
        
        showSubstitutions = concatMap $ uncurry $ printf "%c -> %c\n"

        showFrequencies = concatMap $ uncurry $ printf "%c\t%5.2f%%\n"

        getFrequency :: String -> [(Char, Double)]
        getFrequency txt = [ (c, charFreq c) | c <- ['a'..'z']]
          where filteredTxt = filter (`elem` ['a'..'z']) $ map toLower txt
                l = length filteredTxt
                charFreq c = fromIntegral (length (filter (==c) filteredTxt)) * 100.0 / fromIntegral l

        failEmpty [] = error "No input given, exiting."
        failEmpty _  = return ()