playfair.hs

import Data.List
import Data.Char
import Data.List.HT

type Password = String
type Cipher = [String]

--this should take 5 letters at a time, removing earlier seen letters
--also we should add the remaining letters of the alphabet
createCipherSquare :: Password -> Cipher
createCipherSquare pw = 
  let 
    captialLetters = map toUpper pw
    continousLetters = filterDuplicateLetters captialLetters ++ removeCipherFromAlphabet captialLetters
  in
    partitionToFive (filter (\x -> not $ x == 'J') continousLetters)

partitionToFive :: [Char] -> Cipher
partitionToFive [] = []
partitionToFive xs = take 5 xs : partitionToFive (drop 5 xs)

filterDuplicateLetters :: Password -> [Char]
filterDuplicateLetters pw = nub pw 

removeCipherFromAlphabet :: Password -> [Char]
removeCipherFromAlphabet pw = 
  let alphabet = ['A'..'Z'] in
    [ x | x <- alphabet, not $ elem x pw]

splitIntoPairs :: String -> [String]
splitIntoPairs []       = []
splitIntoPairs [x]      = [[x] ++ "X"]
splitIntoPairs (x:y:xs) = if x == y 
                          then 
                            [x,'X'] : splitIntoPairs xs
                          else
                            [x,y] : splitIntoPairs xs

encode :: String -> Password -> String
encode message pw = 
  let cipher = createCipherSquare pw
      messageInCapitals = (replace "J" "I" . map toUpper . filter isLetter) message
      pairs = splitIntoPairs messageInCapitals
  in
    concat $ map (\x -> codePair cipher 1 x) pairs

decode :: String -> Password -> String
decode message pw =
  let cipher = createCipherSquare pw
      pairs = splitIntoPairs message
  in
    concat $ map (\x -> codePair cipher (-1) x) pairs


applyDownShift :: Cipher -> Int -> (Int, Int) -> (Int, Int) -> String
applyDownShift c dir (r,k) (r1, _) = 
  let newRow =  mod (r + dir) 5 
      newRowPlus = mod (r1 + dir) 5
  in
    getCipherChar c (newRow, k) : [getCipherChar c (newRowPlus, k)]

applyRightShift :: Cipher -> Int -> (Int,Int) -> (Int,Int) -> String
applyRightShift c dir (r,k) (_, k1) = 
  let newCol      = mod (k + dir) 5
      newColPlus  = mod (k1 + dir) 5 
  in
    getCipherChar c (r, newCol) : [getCipherChar c (r, newColPlus)]

applySquareShift :: Cipher -> (Int,Int) -> (Int, Int) -> String
applySquareShift c (rone,kone) (rtwo, ktwo) = 
  let x                 = (rtwo, kone)
      y                 = (rone, ktwo) 
  in
    (getCipherChar c y) : [(getCipherChar c x)]

codePair :: Cipher -> Int -> String -> String
codePair c dir pair = 
  let (rowOne, columnOne) = getRowAndColum c $ head pair
      (rowTwo, columnTwo) = getRowAndColum c $ pair !! 1
  in
    if rowOne == rowTwo 
    then
      applyRightShift c dir (rowOne, columnOne)  (rowOne, columnTwo)
    else if columnOne == columnTwo
    then
      applyDownShift c dir (rowOne,columnTwo) (rowTwo, columnTwo)
    else
      applySquareShift c (rowOne,columnOne) (rowTwo, columnTwo)

findI c []      = -1
findI c (x:xs)  = 
  if c == (fst x) 
  then
    (snd x) 
  else 
    findI c xs

getCipherChar :: Cipher -> (Int, Int) -> Char
getCipherChar c (r,k) = c !! r !! k

getRowAndColum :: Cipher -> Char -> (Int, Int)
getRowAndColum xs c = 
  let zippedRowMap = concat $ map (\x -> zip x [0..]) xs
      zippedColMap = concat $ map (\x -> zip x [0..]) (transpose xs) in
    ( findI c zippedColMap, findI c zippedRowMap)

main :: IO ()
main = do print $ encode "PROGRAMMING PRAXIS" "PLAYFAIR" 
          print $ decode "LIVOBLKZEDOELIYWCN" "PLAYFAIR"