{-# LANGUAGE NoMonomorphismRestriction #-}

# See http://kernelbob.wordpress.com/2011/03/20/same-five-digits/#about for problem
# statement.
import qualified Data.Map as M
import qualified Data.Set as S

squares = map (show . (^2)) [1..]
fiveLong = takeWhile (\ s -> length s < 6) . dropWhile (\ s -> length s < 5)

combinationsOf 0 _ = [[]]
combinationsOf _ [] = []
combinationsOf k (x:xs) = map (x:) (combinationsOf (k-1) xs) ++ combinationsOf k xs

countMap :: [String] -> ([String], M.Map Char Char)
countMap triplet = (triplet, M.map (head . show) . M.fromListWith (+) $ 
                             zip (concat triplet) (repeat 1))

counted n = map countMap . combinationsOf n . fiveLong

fiveChars  = filter (\ (_, m) -> (== 5) $ M.size m)
diffChars  = filter (\ (_, m) -> (== 5) . S.size . S.fromList $ M.elems m)
sameDigits = filter (\ (_, m) -> S.fromList (M.elems m) == S.fromList (M.keys m))
notSelf    = filter (\ (_, m) -> (== 0) . M.size $ M.filterWithKey (==) m)
findUnique = M.filter (\ t -> length t == 3) . flipSingles . getSingles
  where
    getSingles = map (\ (t, m) -> (t, head . M.keys $ M.filter (== '1') m))
    flipSingles = M.fromListWith (++) . map (\ (t, m) -> (m, t))

filterSeq = findUnique . notSelf . sameDigits . diffChars . fiveChars

main = do
  (putStrLn . show . head . M.elems . filterSeq . counted 3) squares