Countdown

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 Op = Add | Sub | Mul | Div 

ops = [Add,Sub,Mul,Div]

instance Show Op where
  show Add = "+"
  show Sub = "-"
  show Mul = "*"
  show Div = "/"

app::Op->(Int->Int->Int)
app Add = (+)
app Sub = (-)
app Mul = (*)
app Div = div

-- evaluate a binary operation, rejecting invalid or redundant cases
apply::Op->Maybe Int->Maybe Int->Maybe Int
apply op mi mj = case (op,mi,mj) of
  (Add, (Just i), (Just j)) | ((i==0) ||(j==0) || (i<j))          -> Nothing
  (Sub, (Just i), (Just j)) | (i<=j)                              -> Nothing
  (Mul, (Just i), (Just j)) | ((i==1) ||(j==1) || (i<j))          -> Nothing
  (Div, (Just i), (Just j)) | ((j<=1) || (i<j) || (mod i j /= 0)) -> Nothing
  _ -> pure (app op) <*> mi <*> mj
  
data Expr = Const Int | E Op Expr Expr (Maybe Int)  

instance Show Expr where
  show (Const x) = show x
  show (E op x y val) = concat ["(",(show x),(show op),(show y),"=",(drop 5 $ show val),")"]
 
eval::Expr->Maybe Int
eval (Const i) = Just i
eval (E _ _ _ val) = val

makeExps::(Expr->Expr->[Expr])
makeExps x y = 
   let x' = eval x
       y'= eval y in 
   if (x'>=y') then [E op x y (apply op x' y') | op<-ops] 
               else [E op y x (apply op y' x') | op<-ops] 

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:makeExps a b
expressions xs =  
  concat $ 
  map (\(a,b) -> concat [makeExps 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 Const xs


data Expr' = Const' Int | Sum' [Expr'] |  Prod [Expr'] |  Diff Expr' Expr' |  Ratio Expr' Expr' 

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

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

-- combine expressions to a Sum' 
(.+)::Expr'->Expr'->Expr'
(.+) (Sum' xs) (Sum' ys) = Sum' $ srt (xs ++ ys)
(.+) x (Sum' ys) = Sum' $ srt (x:ys)
(.+) (Sum' xs) y = Sum' $ srt (y:xs)
(.+) x y = Sum' $ srt [x,y]

-- combine expressions to a Prod 
(.*)::Expr'->Expr'->Expr'
(.*) (Prod xs) (Prod ys) = Prod $ srt (xs ++ ys)
(.*) x (Prod ys) = Prod $ srt (x:ys)
(.*) (Prod xs) y = Prod $ srt (y:xs)
(.*) x y = Prod $ srt  [x,y]

simplify::Expr->Expr'
simplify (Const i) = Sum' [Const' i]
simplify (E Add x y _) = (.+) (simplify x) (simplify y)
simplify (E Mul x y _) = (.*) (simplify x) (simplify y)

simplify (E Sub x y _) =
  let x' = simplify x
      y' = simplify y
  in case (x',y') of
    (Diff xs xs',Diff ys ys') -> Diff (xs .+ ys') (xs' .+ ys)
    (Sum' xs,Diff ys ys') -> Diff ((Sum' xs) .+ (Sum' [ys])) ys' 
    (Diff xs xs',Sum' ys) -> Diff  xs (xs' .+ (Sum' ys)) 
    _ -> Diff x' y'

simplify (E Div x y _) =
  let x' = simplify x
      y' = simplify y
  in case (x',y') of
    (Ratio xs xs',Ratio ys ys') -> Ratio (xs .* ys') (xs' .* ys)
    (Sum' xs,Diff ys ys') -> Ratio ((Prod xs) .* (Prod [ys])) ys' 
    (Diff xs xs',Sum' ys) -> Ratio  xs (xs' .* (Prod ys)) 
    _ -> Ratio x' y'





main::IO ()
main = do
  args<-getArgs
  let ns = read $ args!!0
  let target = read $ args!!1
  mapM_ print $ nub $ map (show.simplify) $ solutions target ns