simulation.hs

{-# LANGUAGE DeriveGeneric, OverloadedStrings, DeriveAnyClass #-}
{-# LANGUAGE TemplateHaskell, NamedFieldPuns, LambdaCase #-}
{-# LANGUAGE RankNTypes, TypeOperators, Strict #-}
{-# LANGUAGE DuplicateRecordFields #-}

{-| For each simulation, every person is randomly assigned a task, and
normally randomly decides how long each task ends up taking for this
particular simulation; it then walks forward in time until the first
person (or people) complete their task, at which point it randomly
assigns new tasks from the remaining, chooses normal random numbers
for how long they'll take, and continues walking forward in time until
the next person finishes a task.
Once it has a total number of days for work being done. It walks
forward in time skipping holidays and weekends to find a terminating
date that it can report to the caller.
-}
module Metadrift.Internal.Simulation where

import qualified Data.Random.Normal as Normal
import Data.Foldable (toList)
import Data.Maybe.Utils (forceMaybe)
import Data.Maybe (catMaybes, isNothing)
import Control.Applicative((<$>))
import Control.DeepSeq (NFData)
import Control.Monad.IO.Class (liftIO)
import qualified Control.Monad.MonteCarlo as MonteCarlo
import Control.Monad.State (StateT, lift, evalStateT, get, put)
import qualified Data.Aeson.TH as AesonEncoder
import qualified Data.Label as Labels
import Data.Label.Monadic (modify, gets, puts)
import qualified Data.List as List
import Data.Sequence ((|>))
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Data.Time.Clock (UTCTime)
import qualified Data.Time.Clock as Clock
import GHC.Generics (Generic)
import qualified Metadrift.Internal.Resources.User.V1 as User
import qualified Metadrift.Internal.Resources.Card.V1 as Card
import qualified Metadrift.Internal.Utils as Utils
import System.Random.TF (newTFGen)
import System.Random.Shuffle (shuffle')
import qualified Metadrift.Internal.Simulation.SimUser as SimUser
import qualified Control.Lens as Lens

data Excluded = Excluded
  { cardName :: Card.Name
  , missingEstimate :: User.Username
  } deriving (Eq, Ord, Show, NFData, Generic)

$(AesonEncoder.deriveJSON Utils.defaultAesonOptions ''Excluded)

data Retired = Retired
  { retiredUser :: SimUser.T
  , retiredAt :: Double
  } deriving (Eq, Ord, Show, NFData, Generic)

data RetiredResult = RetiredResult
  {
    name :: User.Username
  , retiredDate :: UTCTime
  } deriving (Show)

$(AesonEncoder.deriveJSON Utils.defaultAesonOptions ''RetiredResult)

data Result = Result
  { completionDate :: UTCTime
  , reqPercentile :: Double
  , workingDays :: Double
  , excluded :: [Excluded]
  , retired :: [RetiredResult]
  , totalManDays :: Double
  , totalNonWorkingDays :: Double
  , totalPto :: Double
  } deriving (Generic)

$(AesonEncoder.deriveJSON Utils.defaultAesonOptions ''Result)

data SimState = SimState
  { _days :: Double
  , _rawUsers :: Seq.Seq User.T
  , _users :: Seq.Seq SimUser.T
  , _retiredUsers :: Seq.Seq Retired
  , _cards :: Seq.Seq Card.T
  , _excludedCards :: Set.Set Excluded
  , _simulateTimeOff :: Bool
  } deriving (Show, Generic)

$(Labels.mkLabels [''SimState])

data SimulationResult = SimulationResult { daysToComplete :: Double,
                                           allRetiredUsers :: Seq.Seq Retired,
                                           allExcludedCards :: Set.Set Excluded,
                                           simTotalManDays :: Double,
                                           simTotalNonWorkingDays :: Double,
                                           simTotalPto :: Double
                                         } deriving (Eq, Ord, Show, NFData, Generic)

type SimulationState g a = StateT SimState (MonteCarlo.MonteCarlo g) a

experimentCount :: Int
experimentCount = 10000

-- Documentation recomends the count divided by 200 as a starting paint
experimentChunks :: Int
experimentChunks = 50

-- Normally distributed random selection of days
genBetween
  :: (MonteCarlo.RandomGen g)
  => Double -> Double -> SimulationState g Double
genBetween p5 p95 =
  let mean = (p95 + p5) / 2.0
      stddev = (p95 - mean) / p90PercentileFactor
  in do gen <- lift get
        let (newRandom, newGen) = Normal.normal' (mean, stddev) gen
        lift $ put newGen
        return newRandom

-- Pulled from
-- https://en.wikipedia.org/wiki/Normal_distribution#Quantile_function
p90PercentileFactor :: Double
p90PercentileFactor = 1.644853626951

random :: (MonteCarlo.RandomGen g) => SimulationState g Double
random =
  lift MonteCarlo.random

getAdjustmentFactor :: SimUser.T -> Double -> (Double, Double)
getAdjustmentFactor simUser daysToWork =
  let
    nonWorkingTime = Lens.view SimUser.nonWorkingTimeSpread simUser * daysToWork
    pto = Lens.view SimUser.ptoSpread simUser * daysToWork
  in (nonWorkingTime, pto)

-- Normally distributed random selection of days
getElapsedDaysForTask'
  :: SimUser.T -> Double -> SimUser.CardWorkingDays
getElapsedDaysForTask' simUser elapsedDays =
  let
    (nonWorking, pto) = getAdjustmentFactor simUser elapsedDays
  in SimUser.CardWorkingDays{_total= pto + elapsedDays + nonWorking,
                             _nonWorking = nonWorking,
                             _pto = pto}

-- Normally distributed random selection of days
getElapsedDaysForTask
  :: (MonteCarlo.RandomGen g)
  => SimUser.T -> Double -> Double -> SimulationState g SimUser.CardWorkingDays
getElapsedDaysForTask simUser p5 p95 = do
  elapsedDays <- genBetween p5 p95
  timeOffSimulation <- gets simulateTimeOff
  if timeOffSimulation
  then return $ getElapsedDaysForTask' simUser elapsedDays
  else return SimUser.CardWorkingDays{_total= elapsedDays,
                                 _nonWorking = 0.0,
                                 _pto = 0.0}

countOutDaysFromNow :: Double -> IO Clock.UTCTime
countOutDaysFromNow daysToCount = do
  now <- Clock.getCurrentTime
  return $ Utils.countOutDays now daysToCount

getRandomValueFromSequence
  :: (MonteCarlo.RandomGen g)
  => Seq.Seq a -> SimulationState g (Maybe a)
getRandomValueFromSequence lval =
  if Seq.null lval
    then return Nothing
    else do
      index <- lift $ MonteCarlo.randomR (0, Seq.length lval - 1)
      return (Just $ Seq.index lval index)

stepUser :: Double -> SimUser.T -> SimUser.T
stepUser ldays user =
  case SimUser._workInProgress user of
   Just wip ->
     let previousDaysLeft = SimUser._daysLeft wip
         currentDaysLeft = maximum [0, previousDaysLeft - ldays]
     in if currentDaysLeft == 0
        then Lens.over SimUser.workedCards (|> wip) (user {SimUser._workInProgress = Nothing})
        else user {SimUser._workInProgress = Just (wip {SimUser._daysLeft = currentDaysLeft})}
   Nothing -> user

leastStep :: SimulationState g Double
leastStep = do
  localUsers <- gets users
  let daysLeft = SimUser._daysLeft <$> catMaybes (toList (fmap SimUser._workInProgress localUsers))
  let minDays = minimum (0:daysLeft)
  if minDays <= 0
  then return 1
  else return minDays

stepState :: SimulationState g ()
stepState = do
  leastStepInDays <- leastStep
  modify days (+ leastStepInDays)
  modify users (fmap (stepUser leastStepInDays))
  return ()

getWorkDays
  :: (MonteCarlo.RandomGen g)
  => SimUser.T -> Card.T -> SimulationState g (Maybe SimUser.CardWorkingDays)
getWorkDays simUser Card.T {Card.estimates} =
  let username = Lens.view (SimUser.user . User.username) simUser
  in case List.find ((== username) . Card.username) estimates of
    Nothing -> return Nothing
    Just Card.Estimate {Card.range = Card.Range {Card.p5 = rangeP5
                                                ,Card.p95 = rangeP95}} ->
       Just <$> getElapsedDaysForTask simUser rangeP5 rangeP95

addWorkForUser :: User.Username -> SimUser.WorkInProgress ->
            SimUser.T ->  SimUser.T
addWorkForUser username wip potentialWorker =
 let potentialUsername = Lens.view (SimUser.user . User.username) potentialWorker
 in if username == potentialUsername
    then Lens.over SimUser.workInProgress (\_ -> Just wip) potentialWorker
    else potentialWorker

stepWithDays
  :: (MonteCarlo.RandomGen g)
  => SimUser.T -> SimUser.WorkInProgress -> SimulationState g SimulationResult
stepWithDays user wip =
  let username = Lens.view (SimUser.user . User.username) user
  in do
  modify users (fmap (addWorkForUser username wip))
  runSimulation

stepWithCard
  :: (MonteCarlo.RandomGen g)
  => SimUser.T -> Card.T -> SimulationState g SimulationResult
stepWithCard simUser card =
  let newUsername = Lens.view (SimUser.user . User.username) simUser
  in getWorkDays simUser card >>=
     \case
       Just workingDays -> stepWithDays simUser SimUser.WorkInProgress { SimUser._card = card
  , SimUser._daysLeft = SimUser._total workingDays
  , _workingCardDays = workingDays }
       Nothing -> do
         modify
           excludedCards
           (Set.insert
              Excluded
              { cardName = forceMaybe $ Card.name card
              , missingEstimate = newUsername
              })
         stepWithAvailableUser

getNextUnassignedCardInSeq :: User.Username -> Seq.Seq Card.T -> Int -> SimulationState g (Maybe Card.T)
getNextUnassignedCardInSeq userName workingCards index =
  case Utils.lookup index workingCards of
   Just (card @ Card.T {Card.doer = Nothing}) -> do
     puts cards $ Utils.deleteAt index workingCards
     return $ Just card
   Just _ -> getNextUnassignedCardInSeq userName workingCards (index+1)
   Nothing -> return Nothing

getNextAssignedCardInSeq :: User.Username -> Seq.Seq Card.T -> Int -> SimulationState g (Maybe Card.T)
getNextAssignedCardInSeq userName workingCards index =
  case Utils.lookup index workingCards of
    Just (card@ Card.T {Card.doer = Just userName'}) | userName == userName' -> do
      puts cards (Utils.deleteAt index workingCards)
      return $ Just card
    Just _ -> getNextAssignedCardInSeq userName workingCards (index+1)
    Nothing -> getNextUnassignedCardInSeq userName workingCards 0

areAllUsersRetired :: SimulationState g Bool
areAllUsersRetired = do
  allUsers <- gets users
  return $ Seq.length allUsers == 0

getNextCard
  :: SimUser.T -> SimulationState g (Maybe Card.T)
getNextCard user = do
  allCards <- gets cards
  getNextAssignedCardInSeq (Lens.view (SimUser.user . User.username) user) allCards 0

retireUser :: SimUser.T -> SimulationState g ()
retireUser user =
  let userName = Lens.view (SimUser.user . User.username) user
  in do
    ldays <- gets days
    modify users (Seq.filter (\a -> userName /= Lens.view (SimUser.user . User.username) a))
    modify retiredUsers (\allRetiredUsers ->
                           allRetiredUsers |> Retired{ retiredUser = user, retiredAt = ldays })

packageSimulation ::  SimulationState g SimulationResult
packageSimulation = do
  ldays <- gets days
  allRetiredUsers <- gets retiredUsers
  allExcludedCards <- gets excludedCards
  let totals = fmap (SimUser.getAggregateCardWorkingDays . retiredUser) allRetiredUsers
  return SimulationResult { daysToComplete = ldays,
                            allRetiredUsers = allRetiredUsers,
                            allExcludedCards = allExcludedCards,
                            simTotalManDays = sum $ fmap SimUser._total totals,
                            simTotalNonWorkingDays = sum $ fmap SimUser._nonWorking totals,
                            simTotalPto = sum $ fmap SimUser._pto totals }


stepWithAvailableUser
  :: (MonteCarlo.RandomGen g)
  => SimulationState g SimulationResult
stepWithAvailableUser =
  randomlyChooseAnAvailableUser >>= \case
    Nothing -> do
      allUsersAreRetired <- areAllUsersRetired
      if allUsersAreRetired
        then packageSimulation
        else runSimulation
    Just user ->
      getNextCard user >>= \case
        Just card -> stepWithCard user card
        Nothing -> do
          retireUser user
          stepWithAvailableUser

randomlyChooseAnAvailableUser
  :: (MonteCarlo.RandomGen g)
  => SimulationState g (Maybe SimUser.T)
randomlyChooseAnAvailableUser = do
  availableUsers <- Seq.filter (isNothing . SimUser._workInProgress) <$> gets users
  getRandomValueFromSequence availableUsers

shuffleCards :: MonteCarlo.RandomGen g => Seq.Seq Card.T -> SimulationState g (Seq.Seq Card.T)
shuffleCards unshuffledCards = do
  gen <- lift get
  let shuffledCards = shuffle' (toList unshuffledCards) (Seq.length unshuffledCards) gen
  return $ Seq.fromList shuffledCards

runSimulation
  :: (MonteCarlo.RandomGen g)
  => SimulationState g SimulationResult
runSimulation = do
  stepState
  stepWithAvailableUser

generateSimUser ::  (MonteCarlo.RandomGen g)
  => User.T -> SimulationState g SimUser.T
generateSimUser user = do
   dailyTimeOffSpread <- case User._pto user of
                                 Just User.PTO{User.p5, User.p95} -> do
                                    daysOffInTwoMonthPeriod <- genBetween p5 p95
                                    return $ daysOffInTwoMonthPeriod / 40
                                 Nothing -> return 0.0
   nonWorkingTimeSpread <-
      case Lens.view User.nonCardTime user  of
        Just User.NonCardTime {User.p5, User.p95} ->
          genBetween p5 p95
        Nothing ->
          return 0.0
   return SimUser.T {SimUser._user = user,
               SimUser._ptoSpread = dailyTimeOffSpread,
               SimUser._nonWorkingTimeSpread = nonWorkingTimeSpread,
               SimUser._workInProgress = Nothing,
               SimUser._workedCards = Seq.empty }

startSimulation
  :: (MonteCarlo.RandomGen g)
  => SimulationState g SimulationResult
startSimulation = do
  raw <- gets rawUsers
  simUsers <-  mapM generateSimUser raw
  puts users simUsers
  unshuffledCards <- gets cards
  shuffledCards <- -- the shuffle will hang here if it gets an empty list
    if Seq.length unshuffledCards == 0
    then return unshuffledCards
    else shuffleCards unshuffledCards
  puts cards shuffledCards
  runSimulation

convertToProcessable :: [SimulationResult] -> [SimulationResult]
convertToProcessable =
  List.sortOn daysToComplete

retiredToRetiredResult :: Seq.Seq Retired -> IO (Seq.Seq RetiredResult)
retiredToRetiredResult =
  mapM (\a ->
           do
             targetDate <- countOutDaysFromNow (retiredAt a)
             return RetiredResult { name = Lens.view (SimUser.user . User.username) (retiredUser a)
                                  , retiredDate = targetDate })


processResult :: Int -> Double -> [SimulationResult] -> IO Result
processResult resultCount requestedPercentile results' =
  let results = convertToProcessable results'
      index = ceiling (requestedPercentile * fromIntegral resultCount) - 1
      simulationResult = results !! index
  in do targetDate <- countOutDaysFromNow (daysToComplete simulationResult)
        finalRetiredUsers <- retiredToRetiredResult (allRetiredUsers simulationResult)
        return
          Result
          { completionDate = targetDate
          , reqPercentile = requestedPercentile
          , workingDays = daysToComplete simulationResult
          , excluded = Set.toList (allExcludedCards simulationResult)
          , retired = toList finalRetiredUsers
          , totalManDays = simTotalManDays simulationResult
          , totalNonWorkingDays = simTotalNonWorkingDays simulationResult
          , totalPto = simTotalPto simulationResult
          }

-- Runs a set of montecarlo simulations (the number bound by `experimentCount`)
-- then processes the result finding the number of days (and the projected
-- completion date from today).
run :: Double -> Seq.Seq Card.T -> Seq.Seq User.T -> Bool -> IO Result
run requestedPercentile simCards workers timeOffSimulation =
  if Seq.null workers
    then do
      today <- countOutDaysFromNow 0
      return
        Result
        { completionDate = today
        , reqPercentile = requestedPercentile
        , workingDays = 0.0
        , excluded = []
        , retired = []
        , totalManDays = 0.0
        , totalNonWorkingDays = 0.0
        , totalPto = 0.0
        }
    else let initialState =
               SimState
               { _days = 0
               , _rawUsers = workers
               , _users = Seq.empty
               , _cards = simCards
               , _retiredUsers = Seq.empty
               , _excludedCards = Set.empty
               , _simulateTimeOff = timeOffSimulation
               }
         in do
                g <- liftIO newTFGen
                let result =
                     MonteCarlo.experimentP
                       (evalStateT startSimulation initialState)
                       experimentCount
                       experimentChunks
                       g :: [SimulationResult]
processResult experimentCount requestedPercentile result