elevator-sim.hs

{- stack script
--resolver lts-18.26
--package containers
--package random
-}

{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_GHC -Wall #-}

import Control.Monad.State.Strict (State)
import Data.Function ((&))
import Data.List (genericLength, partition)
import Data.Sequence (Seq, ViewL (..))
import System.Random.Stateful (StateGenM (..), StdGen, newStdGen, runStateGen_, uniformM, uniformRM)

import qualified Data.Sequence as Seq


main :: IO ()
main = do
    stdGen <- newStdGen
    let Elevator{transported} = runStateGen_ stdGen $ \_ -> do
            elevator <- initElevator <$> floorGen floorCount
            runSimulation elevator 0
    putStrLn $ "Average wait time: " <> show (average (waitTime <$> transported))
    putStrLn $ "Average ride time: " <> show (average (rideTime <$> transported))


-- | Number of floors in the building
floorCount :: Int
floorCount = 100


-- | Stop the simulation once the number of transported people reaches this threshold
transportedThreshold :: Int
transportedThreshold = 100


runSimulation :: Elevator -> Int -> Gen Elevator
runSimulation elevator currentTime
    | length (transported elevator) >= transportedThreshold = pure elevator
    | otherwise = do
        mWaiter <- waiterGen currentTime
        let newElevator = timeStep mWaiter currentTime elevator
        runSimulation newElevator (currentTime + 1)


timeStep :: Maybe Person -> Int -> Elevator -> Elevator
timeStep mWaiter currentTime Elevator{currentFloor, destinationFloor, requestQueue, waiters, riders, transported} =
    let (exitedOnThisFloor, ridersFromPast) =
            partition (\person -> toFloor person == currentFloor) riders

        (onboardedOnThisFloor, newWaiters) =
            partition
                (\person -> fromFloor person == currentFloor)
                (waiters <> case mWaiter of Nothing -> []; Just w -> [w])

        newRiders = ridersFromPast <> fmap markOnboarded onboardedOnThisFloor
        newTransported = transported <> fmap markTransported exitedOnThisFloor

        markOnboarded p = p{onboardedAt = currentTime}
        markTransported p = p{droppedOffAt = currentTime}

        (newRequestQueue, newDestinationFloor) =
            let requestQueue' =
                    requestQueue
                        -- Remove requests from people who exited
                        & Seq.filter (== currentFloor)
                        -- Add requests from newly onboarded people
                        & (<> Seq.fromList (toFloor <$> onboardedOnThisFloor))
                        -- Add request from new waiter
                        & case mWaiter of Nothing -> id; Just w -> (Seq.|> fromFloor w)
             in case destinationFloor of
                    -- No destination yet, pick a request from the queue
                    Nothing -> case Seq.viewl requestQueue' of
                        EmptyL -> (Seq.empty, Nothing)
                        newFloor :< rest -> (rest, Just newFloor)
                    Just destFloor
                        -- We're at destination, pick new destionation
                        | destFloor == currentFloor -> case Seq.viewl requestQueue' of
                            EmptyL -> (Seq.empty, Nothing)
                            newFloor :< rest -> (rest, Just newFloor)
                        -- not there yet continue without touching request queue
                        | otherwise -> (requestQueue', Just destFloor)

        newCurrentFloor = case newDestinationFloor of
            Nothing -> currentFloor
            Just dest -> currentFloor + signum (dest - currentFloor)
     in Elevator
            { currentFloor = newCurrentFloor
            , destinationFloor = newDestinationFloor
            , requestQueue = newRequestQueue
            , waiters = newWaiters
            , riders = newRiders
            , transported = newTransported
            }


waiterGen :: Int -> Gen (Maybe Person)
waiterGen arrivedAt = do
    -- 50% chance of person arriving at given time
    arrived <- uniformM StateGenM
    if arrived
        then do
            fromFloor <- floorGen floorCount
            toFloor <- differentFloorGen fromFloor
            pure $
                Just
                    Person
                        { arrivedAt = arrivedAt
                        , fromFloor = fromFloor
                        , toFloor = toFloor
                        , onboardedAt = 0
                        , droppedOffAt = 0
                        }
        else pure Nothing


floorGen :: Int -> Gen Int
floorGen maxFloor =
    uniformRM (1, maxFloor) StateGenM


differentFloorGen :: Int -> Gen Int
differentFloorGen a = do
    b <- floorGen floorCount
    if a == b
        then differentFloorGen a
        else pure b


data Elevator = Elevator
    { currentFloor :: Int
    , destinationFloor :: Maybe Int
    , requestQueue :: Seq Int
    , waiters :: [Person]
    , riders :: [Person]
    , transported :: [Person]
    }
    deriving (Show)


data Person = Person
    { arrivedAt :: Int
    , fromFloor :: Int
    , toFloor :: Int
    , onboardedAt :: Int
    , droppedOffAt :: Int
    }
    deriving (Show)


type Gen = State StdGen


initElevator :: Int -> Elevator
initElevator initialFloor =
    Elevator
        { currentFloor = initialFloor
        , destinationFloor = Nothing
        , requestQueue = Seq.empty
        , waiters = []
        , riders = []
        , transported = []
        }


waitTime :: Person -> Int
waitTime Person{arrivedAt, onboardedAt} =
    onboardedAt - arrivedAt


rideTime :: Person -> Int
rideTime Person{onboardedAt, droppedOffAt} =
    droppedOffAt - onboardedAt


average :: [Int] -> Double
average xs = fromIntegral (sum xs) / genericLength xs