import Control.Monad     (guard)
import Data.Array        (Array, assocs, elems, listArray, (!))
import Data.Char         (toUpper)
import Data.List         (nub, (\\))
import Data.List.Split   (chunksOf)
import Data.Maybe        (listToMaybe)
import Data.String.Utils (replace)

type Square a = Array (Int, Int) a


array2D ::
       (Int, Int) 
    -> [e] -> Square e
array2D maxCoord = listArray ((1, 1), maxCoord)





makeTable :: String -> String
makeTable k = nub key ++ (alpha \\ key)
    where
      alpha = ['A' .. 'Z'] \\ "J"
      key = map toUpper =<< words k


makeSquare :: [a] -> Square a
makeSquare = array2D (5, 5)


showSquare :: Square Char -> String
showSquare d = unlines $ chunksOf 5 (elems d)


getIndex':: (Eq a) => a -> [((Int, Int), a)] -> Maybe (Int, Int)
getIndex' el = fmap fst . listToMaybe . filter ((== el) . snd)

encodePair, decodePair :: Eq a => Square a -> (a, a) -> Maybe (a, a)
encodePair = pairHelper (\x -> if x == 5 then 1 else x + 1)
decodePair = pairHelper (\x -> if x == 1 then 5 else x - 1)

pairHelper :: (Eq t)
    => (Int -> Int) 
    -> Square t 
    -> (t, t) 
    -> Maybe (t, t) 
pairHelper adjust sqr (c1, c2) =
    do let ps = assocs sqr
       
       (x1, y1) <- getIndex' c1 ps
       (x2, y2) <- getIndex' c2 ps
       
       
       guard $ c1 /= c2
       
       let get x = sqr ! x
       
       Just $
           
           case () of
             () | y1 == y2 ->
                    (get (adjust x1, y1), get (adjust x2, y2))
                | x1 == x2 ->
                    (get (x1, adjust y1), get (x2, adjust y2))
                | otherwise ->
                    (get (x1, y2), get (x2, y1))


parsePair :: String -> [(Char, Char)]
parsePair = fmap (\[x, y] -> (x, y)) . words . fmap toUpper


unparsePair :: [(Char, Char)] -> String
unparsePair = unwords . fmap (\(x, y) -> [x, y])

codeHelper :: (Square Char -> (Char, Char) -> Maybe (Char, Char))
    -> String -> String -> Maybe String
codeHelper subs key =
    fmap unparsePair .
    mapM (subs (makeSquare $ makeTable key)) .
    parsePair

playfair, unplayfair :: String -> String -> Maybe String
playfair key = codeHelper encodePair key . formatEncode
unplayfair = codeHelper decodePair

formatEncode :: String -> String
formatEncode =
    map toUpper .
    unwords .
    map (\[x, y] -> if x == y then [x, 'x'] else [x, y]) .
    chunksOf 2 .
    replace "j" "i" .
    concatMap adjustLength .
    words .
    filter (\n -> n `elem` (['A'..'Z'] ++ ['a'..'z']))
    where
      adjustLength str
          | odd (length str) = str ++ "x"
          | otherwise = str