module Main where

import Data.List

data NestedList a = Elem a | List [NestedList a] deriving Show
type NestedInt = NestedList Int

test :: (Show a, Eq a) => a -> a -> String
test x y = if x == y then "Pass" else "expected " ++ show x ++ " but got " ++ show y

primes n = sieve [2..n]
  where sieve (x:xs) = x : sieve (xs \\ [x,x+x..n])
        sieve [] = []

primeFactors :: Int -> [Int]
primeFactors n
  | null factors = [n]
  | otherwise   = let f = head factors in f:primeFactors ( n `div` f)
  where factors = filter (\x -> n `rem` x == 0) $ primes (n-1)

encode' :: NestedInt -> NestedInt
encode' (List xs) = List $ map encode' xs
encode' (Elem n)
  | n == 2    = List []
  | otherwise = List $ map (encode' . Elem) $ primeFactors n'
  where primes' = primes n
        n'      = if n `elem` primes' then pIndex else n
        pIndex  = fst $ head $ filter ((n==) . snd) $ zip [1..] primes'

encode :: Int -> String
encode = filter (`elem` "[]") . show . (read :: String -> NestedInt) . show . encode' . Elem

main :: IO ()
main = do
  print $ test "[[[]][[]]]" $ encode 9
  print $ test "[[][[[]][[]]]]" $ encode 46