Main.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import           Control.Applicative
import           Control.Monad.ST
import           Data.Bool
import qualified Data.Foldable as F
import           Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Algorithms.Radix as VA
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Unboxed.Mutable as VM
import           GHC.Exts (inline)
import           Linear (V3(..))
import           System.Environment (getArgs)

type Forest = V3 Int

meals :: Vector Forest
meals = V.fromList [ V3 (-1) (-1)   1
                   , V3 (-1)   1  (-1)
                   , V3   1  (-1) (-1)
                   ]

nubVector :: (Ord a, VA.Radix a, VM.Unbox a) => Vector a -> Vector a
nubVector iv = runST $ do
    v <- V.thaw iv
    inline $ VA.sort v
    V.unsafeFreeze =<< uniq v

uniq :: (Ord a, VM.Unbox a) => VM.MVector s a -> ST s (VM.MVector s a)
uniq v = if VM.null v then return v else go 1 1
  where go i j | j >= VM.length v = return $ VM.take i v
               | otherwise        = (==) <$> VM.read v (j - 1) <*> VM.read v j >>= bool ((VM.read v j >>= VM.write v i) >> go (i + 1) (j + 1)) (go i (j + 1))

solve :: Vector Forest -> Vector Forest
solve fs = if V.null sol then solve fs' else sol
  where fs' = nubVector $ flip V.concatMap fs $ \f -> V.filter (F.all (>= 0)) $ V.map (f +) meals
        sol = V.filter ((>= 2) . length . filter (== 0) . F.toList) fs'

main :: IO ()
main = do
    [g, w, l] <- map read <$> getArgs
    print $ solve $ V.singleton (V3 g w l)

instance (VA.Radix e) => VA.Radix (V3 e) where
    passes _ = 3 * VA.passes (undefined :: e)
    {-# INLINE passes #-}
    size _ = VA.size (undefined :: e)
    {-# INLINE size #-}
    radix n (V3 x y z) | n <     p = VA.radix  n          x
                       | n < 2 * p = VA.radix (n -     p) y
                       | otherwise = VA.radix (n - 2 * p) z
      where p = VA.passes (undefined :: e)
    {-# INLINE radix #-}