Countdown 2

import System.Environment
import Data.List (partition,nub,sortBy)
import Data.Maybe (isJust)
import qualified Data.IntMap.Strict as M
import Control.Monad (filterM)

data Expr = Single Int | 
            Sum' [Expr] (Maybe Int) | 
            Prod [Expr] (Maybe Int) | 
            Diff Expr Expr (Maybe Int) | 
            Ratio Expr Expr (Maybe Int)

-- sort the arguments of Sum' or Prod - the ordering only needs to be consistent
srt::[Expr]->[Expr]
srt xs =  sortBy (\x y -> compare (show (eval x)) (show (eval y))) xs

instance Show Expr where
  show (Single i) = concat [(show i)," "]
  show (Sum' xs val) = concat ["(Sum [",(concat (map show (srt xs))),"])"]
  show (Prod xs val) = concat ["(Prod [",(concat (map show (srt xs))),"])"]
  show (Diff x y val) = concat ["(",(show x), "-", (show y),")"]
  show (Ratio x y val) = concat ["(",(show x), "/", (show y),")"]

-- val is maintained by the combinators
eval::Expr->Maybe Int
eval (Single i) = Just i
eval (Sum' _ val) = val
eval (Prod _ val) = val
eval (Diff _ _ val) = val
eval (Ratio _ _ val) = val

(~+)::Maybe Int->Maybe Int->Maybe Int
(~+) x y = ((pure (+))<*>x<*>y)

(~-)::Maybe Int->Maybe Int->Maybe Int
(~-) x y = (pure (-))<*>x<*>y

(~*)::Maybe Int->Maybe Int->Maybe Int
(~*) x y = (pure (*))<*>x<*>y

-- i>=j enforced at call site (combine)
(~/)::Maybe Int->Maybe Int->Maybe Int
(~/) (Just i) (Just j) | ((j>1) && ((mod i j)==0)) = Just (div i j)
(~/) _ _ = Nothing

-- cases: Sum' Sum', Sum' Diff, Diff Sum', Diff Diff, other Sum', Sum' other, other other
(.+)::Expr->Expr->Expr
(.+) (Sum' xs valx)   (Sum' ys valy)   = Sum' (xs ++ ys) ((~+) valx valy)
(.+) (Sum' xs valx)   (Diff y y' valy) = Diff ((.+) (Sum' xs valx) y)  ((.+) y' (Sum' [] (Just 0))) ((~+) valx valy)   
(.+) (Diff y y' valy) (Sum' xs valx)   = (.+) (Sum' xs valx)  (Diff y y' valy) -- commutative
(.+) (Diff x x' valx) (Diff y y' valy) = Diff ((.+) (Sum' [x] (eval x)) (Sum' [y'] (eval y')))  ((.+) (Sum' [x'] (eval x')) (Sum' [y] (eval y))) ((~+) valx valy)
(.+) (Sum' xs valx) y = Sum' (y:xs) ((~+) valx (eval y))
(.+) y (Sum' xs valx) = (.+) (Sum' xs valx) y -- commutative
(.+) x y = Sum' [x,y] ((~+) (eval x) (eval y))

-- cases: Prod Prod, Prod Ratio, Ratio Prod, Ratio Ratio, other Prod, Prod other, other other
(.*)::Expr->Expr->Expr
(.*) (Prod xs valx)    (Prod ys valy)    = Prod (xs ++ ys) ((~*) valx valy)
(.*) (Prod xs valx)    (Diff y y' valy)  = Diff ((.*) (Prod xs valx) y)  ((.*) y' (Prod [] (Just 0))) ((~*) valx valy)   
(.*) (Ratio y y' valy) (Prod xs valx)    = (.*) (Prod xs valx)  (Ratio y y' valy) -- commutative
(.*) (Ratio x x' valx) (Ratio y y' valy) = Ratio ((.*) (Prod [x] (eval x))  (Prod [y'] (eval y')))   ((.*) (Prod [x'] (eval x')) (Prod [y] (eval y))) ((~*) valx valy)
(.*) (Prod xs valx) y = Prod (y:xs) ((~*) valx (eval y))
(.*) y (Prod xs valx) = (.*) (Prod xs valx) y -- commutative
(.*) x y = Prod [x,y] ((~*) (eval x) (eval y))

-- cases: Diff Diff, Diff Sum', Sum' Diff, other Diff, Diff other, other other
(.-)::Expr->Expr->Expr
(.-) (Diff x x' valx) (Diff y y' valy) = Diff ((.+) x y') ((.+) x' y) ((~-) valx valy)
(.-) (Diff y y' valy) (Sum' xs valx)   = Diff y ((.+) y' (Sum' xs valx)) ((~-) valy valx) 
(.-) (Sum' xs valx)   (Diff y y' valy) = Diff ((.+) y (Sum' xs valx)) y' ((~-) valx valy) 
(.-) (Diff x x' valx) y = Diff x ((.+) y x') ((~-) valx (eval y)) 
(.-) x (Diff y y' valy) = Diff ((.+) x y') y ((~-) (eval x) valy) 
(.-) x y = Diff x y ((~-) (eval x) (eval y))

-- cases: Ratio Ratio, Ratio Sum', Sum' Ratio, other Ratio, Ratio other, other other
(./)::Expr->Expr->Expr 
(./) (Ratio x x' valx) (Ratio y y' valy) = Ratio ((.*) x y') ((.*) x' y) ((~/) valx valy)
(./) (Ratio y y' valy) (Prod xs valx)   = Ratio y ((.*) y' (Prod xs valx)) ((~/) valy valx) 
(./) (Prod xs valx)   (Ratio y y' valy) = Ratio ((.*) y (Prod xs valx)) y' ((~/) valx valy) 
(./) (Ratio x x' valx) y = Ratio x ((.*) y x') ((~/) valx (eval y)) 
(./) x (Ratio y y' valy) = Ratio ((.*) x y') y ((~/) (eval x) valy) 
(./) x y = Ratio x y ((~/) (eval x) (eval y))

combine::Expr->Expr->[Expr]
combine  x y = 
   let x' = eval x
       y'= eval y in 
  filter (isJust.eval) $ 
  case (compare x' y') of
    GT -> [(.+),(.-),(.*),(./)]<*>[x]<*>[y]
    EQ -> [(.+),(.*),(./)]<*>[x]<*>[y] -- no subtraction because 0 would be redundant
    LT -> [(.+),(.-),(.*),(./)]<*>[y]<*>[x]
    
check::Int->Expr->Bool
check target exp = case (eval exp) of
  Just i -> i == target
  _      -> False 

-- combinatorial functions
-- non-trivial masks: each generated partition must be non-empty to ensure recursion progresses
-- 2 and 1 are special cases at call site in expressions
-- don't need reciprocal masks: TTF will generate the same partition as FFT
masks::Int->[[Bool]]
masks 1 = [[True]]  
masks 2 = [[True,True],[False,True]]
masks n = filter (any (==False)) [x : y | x<-[True,False], y<-masks (n-1)]

-- memoised; hardcoded max length 6
masksM::Int->[[Bool]]
masksM i = let m = M.fromAscList $ map (\x->(x,masks x)) [1..6]
  in (M.!) m i

partitions::[a]->[([a],[a])]
partitions xs = 
  map (\(x,y)->((map fst x,map fst y))) $ 
  map (partition (\x->(snd x)==True)) $ 
  map (zip xs) $  
  masksM (length xs) 

powerset::[a]->[[a]]
powerset xs = filterM (\x->[True,False]) xs 

-- all the expressions that use all the numbers in the partition
expressions::[Expr]->[Expr]
expressions [] = []
expressions (a:[]) = [a]
expressions (a:b:[]) = a:b:combine a b
expressions xs =  
  concat $ 
  map (\(a,b) -> concat [combine x y | x<-a,y<-b]) $ -- all the expressions combining the partitions
  map (\(a,b)->(expressions a,expressions b)) $       -- all the expressions from each partition 
  partitions xs                                       -- nontrivial ways to split the list 

solutions'::Int->[Int]->[Expr]
solutions' target xs =
  filter (check target) $ 
  concat $ 
  map expressions $ 
  powerset $ 
  map Single xs

-- deduplicate
solutions::Int->[Int]->[String]
solutions target ns = nub $ map show $ solutions' target ns

main::IO ()
main = do
  args<-getArgs
  let ns = read $ args!!0
  let target = read $ args!!1
  mapM_ putStrLn $ solutions target ns