import Data.Bits
import Data.List
import Data.List.Split

type Eqn = (Int, [Int])
data Op = Add | Mul
data Op2 = Add2 | Mul2 | Cat2 deriving Show

atoi :: String -> Int
atoi = read

process :: String -> Eqn
process line = 
  let [targs, inss] = splitOn ": " line
      targ = atoi targs
      ins = map atoi $ words inss in
        (targ, ins)

opList :: [Op] -> Int -> Int -> [Op]
opList acc 0 _ = acc
opList acc w n =
  opList ((case bit of {0 -> Add; otherwise -> Mul}):acc) dw n
  where dw = w - 1
        bit = n .&. (1 `shift` dw)

reifies :: Eqn -> [Op] -> Bool
reifies (targ, (h:t)) ops =
  foldl operate h tasks == targ
  where tasks = zip ops t
        operate = \acc (op, arg) -> case op of {Add -> acc + arg; Mul -> acc * arg}

reifiable :: Eqn -> Bool
reifiable eqn@(_, args) =
  any (reifies eqn) $ map (opList [] l) [0..(2^l - 1)]
  where l = (length args) - 1

op2List :: [Op2] -> Int -> Int -> [Op2]
op2List acc 0 _ = acc
op2List acc w n =
  op2List ((case trit of {0 -> Add2; 1 -> Mul2; 2 -> Cat2}):acc) dw n
  where dw = w - 1
        trit = (n `div` (3 ^ dw)) `mod` 3

reifies2 :: Eqn -> [Op2] -> Bool
reifies2 (targ, (h:t)) ops =
  foldl operate h tasks == targ
  where tasks = zip ops t
        operate = \acc (op, arg) -> case op of
                                      Add2 -> acc + arg
                                      Mul2 -> acc * arg
                                      Cat2 -> atoi (show acc ++ show arg)

reifiable2 :: Eqn -> Bool
reifiable2 eqn@(_, args) =
  any (reifies2 eqn) $ map (op2List [] l) [0..(3^l - 1)]
  where l = (length args) - 1

part1 :: [Eqn] -> Int
part1 eqns =
  foldl sumCar 0 $ filter reifiable eqns
  where sumCar = \acc (i, _) -> acc + i

part2 :: [Eqn] -> Int
part2 eqns =
  foldl sumCar 0 $ filter reifiable2 eqns
  where sumCar = \acc (i, _) -> acc + i

main = do
  file <- readFile "day07.input"
  let lns = lines file
      eqns = map process lns
  putStrLn ("Part 1: " ++ (show $ part1 eqns))
  putStrLn ("Part 2: " ++ (show $ part2 eqns))