Solving Work plan by SMT solver

planning.hs

{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE DeriveFoldable, GeneralizedNewtypeDeriving, OverloadedStrings #-}
{-# LANGUAGE RecordWildCards, TemplateHaskell                              #-}
module Main where
import           Control.Arrow               (second)
import           Control.Lens                hiding ((.>))
import           Control.Monad
import qualified Data.Foldable               as F
import           Data.Function
import           Data.List
import           Data.List.Split             (splitOn)
import           Data.Maybe
import           Data.Monoid
import           Data.Ord
import           Data.SBV
import           Data.SBV.Internals
import           Data.String
import           Data.Time
import           Data.Time.Calendar.WeekDate

newtype Job = Job { getJob :: String }
            deriving (IsString, Show, Eq, Ord)

data DayOfWeek = Sun | Mon | Tue | Wed | Thr | Fri | Sat
               deriving (Read, Show, Eq, Ord, Enum)

isWeekDay :: DayOfWeek -> Bool
isWeekDay Sun = False
isWeekDay Sat = False
isWeekDay _   = True

data DaySpec = Once Day
             | Weekly DayOfWeek
             | Daily (Maybe Day) (Maybe Day)
             deriving (Read, Show, Eq, Ord)

data Period = PerDay | PerWeek
            deriving (Read, Show, Eq, Ord, Enum)

data Constraint' job =
  TotalJobTime
  { job    :: job
  , duration :: DiffTime
  }
  | Scheduled
    { job     :: job
    , daySpec :: DaySpec
    , begins  :: TimeOfDay
    , ends    :: TimeOfDay
    }
  | Busy
    { label   :: String
    , daySpec :: DaySpec
    , begins  :: TimeOfDay
    , ends    :: TimeOfDay
    }
  | Holiday
    { label   :: String
    , daySpec :: DaySpec
    }
  | UpperBound
    { period   :: Period
    , duration :: DiffTime
    }
  | OtherConstraint (Day -> TimeInterval -> Job -> SBool)
  deriving (Foldable)

type Constraint = Constraint' Job

jobs :: Constraint -> [Job]
jobs = nub . sort . F.toList

data Config = Config { openingTime  :: TimeOfDay
                     , closingTime  :: TimeOfDay
                     , periods      :: DiffTime
                     , firstDay     :: Day
                     , lastDay      :: Day
                     , constraints  :: [Constraint]
                     }
            deriving ()

seconds :: Integer -> DiffTime
seconds = fromInteger

minutes :: Integer -> DiffTime
minutes = (60*) . fromInteger

hours :: Integer -> DiffTime
hours = (60*) . minutes

days :: Integer -> DiffTime
days = (24*) . hours

tod :: Int -> Int -> TimeOfDay
tod hh mm = TimeOfDay hh mm 00

day :: Integer -> Int -> Int -> Day
day = fromGregorian

dayOfWeek :: Day -> DayOfWeek
dayOfWeek = toEnum . (`mod` 7) . view _3 . toWeekDate

genDays :: Config -> [(Day, [(TimeOfDay, TimeOfDay)], DayOfWeek)]
genDays conf@Config{..} = [ (d, divBlocks d, dayOfWeek d)
                     | d <- [firstDay .. lastDay]
                     , all (not . (`matchesDay` d)) holidays]
  where
    holidays = [dspec | Holiday _ dspec <- constraints]
    divBlocks d =
      let assigned = [(s, e) | Scheduled _ dspec s e <- constraints
                             , dspec `matchesDay` d
                             ]
          unavail  = unavailables conf d
          rest     = F.foldrM subtInterval (openingTime, closingTime) $
                     assigned ++ map snd unavail
      in sortBy (comparing fst) $ concatMap (divideByPeriod periods) rest ++ assigned

type TimeInterval = (TimeOfDay, TimeOfDay)

divideByPeriod :: DiffTime -> TimeInterval -> [TimeInterval]
divideByPeriod p (s, e) =
  let (ds, de) = (timeOfDayToTime s, timeOfDayToTime e)
  in [ (ds + fromInteger n * p, min de (ds + fromInteger (n + 1) * p))
       &  both %~ timeToTimeOfDay
     | n <- [0.. ceiling ((de - ds) / p) - 1
            ]
     ]

subtInterval :: Ord a
             => (a,a)
             -> (a,a)
             -> [(a,a)]
subtInterval d@(delS, delE) t@(targS, targE)
  | d `overlaps` t = filter (not . uncurry (>=)) $
                     [(targS, min delS targE) , (min delE targE, targE)]
  | otherwise      = [t]

matchesDay :: DaySpec -> Day -> Bool
matchesDay (Once e)    d = d == e
matchesDay (Weekly w)  d = dayOfWeek d == w
matchesDay (Daily s e) d =
  fromMaybe (day 0 1 1) s <= d && d <= fromMaybe (day 99999 12 31) e

unavailables :: Config -> Day
             -> [(String,  (TimeOfDay, TimeOfDay))]
unavailables Config{..} d =
  [ (lab, (s, e))
  | Busy lab spc s e <- constraints
  , spc `matchesDay` d
  ]

workName :: Day -> TimeInterval -> Job -> String
workName d (s, e) (Job j) = intercalate "-" [showDate d, show s, show e, j]
  where
    showDate = formatTime defaultTimeLocale "%Y%m%d"

intervalSize :: TimeInterval -> DiffTime
intervalSize (s, e) = timeOfDayToTime e - timeOfDayToTime s

problem :: Config -> Symbolic SBool
problem conf@Config{..} = do
  let bs = genDays conf
      js = sort $ nub $ concatMap jobs constraints
      unConstr = [ (j, [(dspec, (s, e))
                       | Scheduled j' dspec s e <- constraints, j' == j])
                 | j <- js \\ [j | TotalJobTime j _ <- constraints ]
                 ]
      cs = [c | OtherConstraint c <- constraints ]
  workJobs <- forM bs $ \(d, ps, w) -> do
    j <- sWord16 $ intercalate "-" ["job", show d]
    constrain $ 0 .<= j &&& j .<= fromIntegral (length js - 1)
    vars <- forM ps $ \ p@(s,e) ->  do
      wks <- forM js $ \jb -> do
        v <- sBool (workName d p jb)
        F.for_ (lookup jb unConstr) $ \drs ->
          unless (any (\(ds, ran) -> ds `matchesDay` d && ran == p) drs) $
          constrain $ bnot v
        constrain $ bOr [ bnot (c d p jb) | c <- cs] ==> bnot v
        return v
      constrain $
        bAnd $
          -- Only one job should be filled for one block.
          [ wk ==> j .== fromInteger n
          | (n, wk) <- zip [0..] wks
          ] ++
          -- Process scheduled jobs
          [ wks !! n
          | Scheduled lab d' s' e' <- constraints
          , s' == s, e' == e, d' `matchesDay` d
          , n <- maybeToList $ findIndex (== lab) js
          ]
          -- If no upper bound specified and unscheduled,
          -- then it must be zero.
      return (p, wks)
    return (d, vars, w)
  -- Processing per-day upper bound
  forM_ workJobs $ \(_, vs, _) -> do
    let ubs = [ub | UpperBound PerDay ub <- constraints ]
        wub = floor $ minimum (hours 24 : ubs)
        total = getSum $ foldMap calcDur vs
    constrain $ total .<= fromInteger wub
  -- Processing per-week upper bound
  forM_ (groupBy ((<) `on` view _3) workJobs) $ \dpsws -> do
    let ubs = [ub | UpperBound PerWeek ub <- constraints ]
        wub = floor $ minimum (hours 24 : ubs)
        total = getSum $ foldMap (foldMap calcDur . view _2) dpsws
    constrain $ total .<= fromInteger wub
    return ()

  -- Requiring total job time constraint
  let slots = map (getSum . foldMap (Sum . calcDur . second pure)) $
              foldr (zipWith (++)) (replicate (length js) []) $
              filter (not . null) $
              map (transpose . map (uncurry (map . (,))) . view _2) $
              workJobs
  forM_ (zip js slots) $ \ (job, dur) -> do
    let ubs = [ floor d | TotalJobTime lab d <- constraints, job == lab]
    F.for_ (listToMaybe ubs) $ \t -> constrain $ (getSum dur) .== fromInteger t
    return ()
  solve []

calcDur :: (TimeInterval, [SBool]) -> Sum SWord32
calcDur (p, wks) = Sum $ fromIntegral (floor $ intervalSize p) * oneIf (bOr wks)

data Work = Work { jobName   :: Job
                 , workDay   :: Day
                 , workDur   :: TimeInterval
                 , workBreak :: DiffTime
                 }
          deriving (Show, Eq, Ord)

isZero :: Work -> Bool
isZero w = fst (workDur w) == snd (workDur w)

instance Monoid Work where
  mempty = Work "" (day 0 1 1) (tod 0 0, tod 0 0) 0
  l@(Work j d (s, e) _) `mappend` r@(Work _ _ (s', e') _)
   | isZero l = r
   | isZero r = l
   | otherwise =
    let dur = calcWorkDur l + calcWorkDur r
        st  = min s s'
        en  = max e e'
        brk = timeOfDayToTime en - timeOfDayToTime st - dur
    in Work j d (st, en) brk

calcWorkDur :: Work -> DiffTime
calcWorkDur (Work _ _ (s, e) b) = timeOfDayToTime e - timeOfDayToTime s - b

{-@ disjoint :: Ord a
             => { (s, e) : (a, a) | s <= e}
             -> { (s, e) : (a, a) | s <= e}
             -> Bool
@-}
disjoint :: Ord a => (a, a) -> (a, a) -> Bool
disjoint (s, e) (s', e') = e < s' || e' < s

overlaps :: Ord a => (a, a) -> (a, a) -> Bool
overlaps l r = not $ disjoint l r

decodeName :: String -> Maybe (Day, TimeInterval, Job)
decodeName var = do
  (x : y : z : rest) <- return $ splitOn "-" var
  d <- parseTimeM True defaultTimeLocale "%Y%m%d" x
  (s, "") <- listToMaybe $ reads y
  (e, "") <- listToMaybe $ reads z
  return $ (d, (s, e), Job $ intercalate "-" rest)

decodeSMTModel :: Config -> SMTModel -> [Work]
decodeSMTModel Config{..} (SMTModel dic) =
  [ Work j d (s, e) 0
  | (lab, cw@(CW KBool _)) <- dic
  , (d, (s, e), j) <- maybeToList $ decodeName lab
  , cwToBool cw
  ]

decodeSatResult :: Config -> SatResult -> [Work]
decodeSatResult c (SatResult (Satisfiable _ m)) = decodeSMTModel c m
decodeSatResult _ _ = []

formatWork :: Work -> String
formatWork Work{..} =
  intercalate "\t" [ show workDay , show (fst workDur)
                   , show (snd workDur)
                   , show (floor workBreak) ++ "秒"
                   , getJob jobName
                   ]

unionSucc :: [Work] -> [Work]
unionSucc =
  map mconcat . groupBy ((==) `on` workDay) . sortBy (comparing workDay <> comparing (fst . workDur))

main :: IO ()
main = mapM_ (putStrLn . formatWork) . unionSucc . decodeSatResult myConfig =<< sat (problem myConfig)

myConfig :: Config
myConfig = Config { openingTime = tod 10 00
                  , closingTime = tod 19 00
                  , periods     = minutes 60
                  , firstDay    = day 2016 10 01
                  , lastDay     = day 2017 1  31
                  , constraints = constrs
                  }
  where
    constrs = [ UpperBound PerDay  $ hours 8
              , UpperBound PerWeek $ hours 35
              , Busy "Seminar" (Weekly Tue) (tod 15 00) (tod 17 00)
              , Busy "Break"   (During Daily Nothing Nothing) (tod 12 15) (tod 13 15)
              , Busy "Workshop"
                  (During Daily (Just $ day 2016 11 15) (Just $ day 2016 11 15))
                  (tod 00 00) (tod 23 59)
              , TotalJobTime "Server Maintainance"  (hours 25)
              , TotalJobTime "Teaching Assistant A"  (hours 18)
              , Holiday "Taiku" $ Once $ day 2016 10 10
              , Holiday "Bunka" $ Once $ day 2016 11 03
              , Holiday "Kinro" $ Once $ day 2016 11 23
              , Holiday "Tenno" $ Once $ day 2016 12 23
              , Holiday "Sat"   $ Weekly Sat
              , Holiday "Sun"   $ Weekly Sun
              , Holiday "New Year" $ During Daily (Just $ day 2017 01 01) (Just $ day 2017 01 03)
              , Scheduled
                  "Teaching Assistant A"
                  (During (Weekly Mon) Nothing (Just $ day 2016 11 31))
                  (tod 13 00)  (tod 14 30)
              ]