Haskell Quiz/The Solitaire Cipher/Solution Thiago Arrais

module Main where

import Char(chr, isAlpha, ord, toUpper)
import System.Environment(getArgs, getProgName)

toUpperCase = map toUpper
toLetter n = chr $ ord 'A' + (n - 1) `mod` 26
toNumber l = ord l - ord 'A' + 1

split5 cs
    | length cs > 5 = (take 5 cs) : split5 (drop 5 cs)
    | otherwise     = [cs]

fill cs = cs ++ replicate (5 - length cs) 'X'

-- Filters alpha characters and splits them into groups of five
sanitize:: String -> [String]
sanitize cs = reverse $ (fill.head) rchunks : tail rchunks
    where rchunks = reverse.split5.filter isAlpha.toUpperCase $ cs

unkeyeddeck :: [Int]
unkeyeddeck = [1..54]

jokerA = 53
jokerB = 54
isJoker = (`elem` [jokerA, jokerB])

-- Pushes a card (j) down once
push' j xs = if not (null right)
             then left ++ head right : j : tail right
             else head left : j : tail left
    where (left,_:right) = break (== j) xs

-- Pushes a card (j) down a given number (n) of times
push j n = (!! n) . iterate (push' j)

pushJokerA = push jokerA 1
pushJokerB = push jokerB 2

-- Performs a triplecut around the first two cards that satisfy a predicate (p)
tripleCut p xs = bottom ++ j1 : (middle ++ j2 : top)
    where (top,j1:b1) = break p xs
          (middle,j2:bottom) = break p b1

countCut n xs = (reverse.tail $ rbottom) ++ top ++ [head rbottom]
    where top = take n xs
          rbottom = reverse.drop n $ xs

-- Performs a count cut by the number written on the bottom card
deckCut xs = countCut (last xs) xs

valueFor 54 = 53 -- B joker's value is 53
valueFor n = n

-- Shuffles the deck once
nextDeck = deckCut.tripleCut isJoker.pushJokerB.pushJokerA

-- Shuffles the deck once and extracts the resulting letter
stepStream :: (String, [Int]) -> (String, [Int])
stepStream (_, oldDeck) = (letter $ number newDeck, newDeck)
    where newDeck = nextDeck oldDeck
          number deck@(n:_) = deck !! valueFor n
          letter n = if isJoker n then "" else [toLetter n]

-- The keystream generated by an unkeyed deck
keystream = concat [c | (c,_) <- tail $ iterate stepStream ([], unkeyeddeck)]

-- Combines an input string (xs) and the default keystream by applying the
-- given operation (f). This is the function that does the encoding/decoding
codeWith f xs = unwords.sanitize.letterize $
                zipWith f (numberize letters) (numberize keyletters)
    where keyletters = take (length letters) keystream
          numberize = map toNumber
          letterize = map toLetter
          letters = concat $ sanitize xs

encode, decode :: String -> String
encode = codeWith (+)
decode = codeWith (-)

-- An action that applies the coding function (f) to a set of words
-- and prints the resulting code
printCode f = putStrLn . f . unwords

main = do args <- getArgs
          case args of
              ("d":ws@(_:_)) -> printCode decode ws
              ("e":ws@(_:_)) -> printCode encode ws
              _              -> getProgName >>=
                  \n -> putStrLn $ "Usage: " ++ n ++ " <d/e> <phrase>"