{-# LANGUAGE NamedFieldPuns #-}
module Timeline (
      State(..)
    , stateAt
  ) where

import Data.List (sortOn)
import Data.Map ((!), Map, mapWithKey)
import Data.Time (Day, diffDays)
import Event (Event(..), EventType(..))
import Medicine (MedicineName, Pharmacy)
import YAML ((.:), Value(..), YAML(..))

data State = State {
      day :: Day
    , stock :: Map MedicineName Float
    , consumptionRate :: Map String Float
  }

instance YAML State where
  toYAML (State {day, stock, consumptionRate}) = Object [
        "date" .: day
      , "stock" .: stock
      , "consumption rates" .: consumptionRate
    ]

initState :: Pharmacy -> State
initState pharmacy = State {
      day = toEnum 0
    , stock = const 0 <$> pharmacy
    , consumptionRate = const 0 <$> pharmacy
  }

applyEvent :: State -> Event -> State
applyEvent state (Event {date, amounts, eventType}) =
  case eventType of
    Prescription -> newState {consumptionRate = amounts}
    Provisioning -> newState {stock = mapWithKey addAmount $ stock newState}
  where
    newState = setDay date state
    addAmount medicineName = (+ (amounts ! medicineName)) 

setDay :: Day -> State -> State
setDay newDay state@(State {day, stock, consumptionRate}) = state {
      day = newDay
    , stock = mapWithKey consume stock
  }
  where
    duration = toEnum . fromEnum $ newDay `diffDays` day
    consume medicineName initialAmount =
      max 0 $ initialAmount - duration * (consumptionRate ! medicineName)

stateAt :: Day -> Pharmacy -> [Event] -> State
stateAt targetDay pharmacy =
  setDay targetDay . lastState
  where
    lastState =
        foldl applyEvent (initState pharmacy)
      . sortOn date
      . filter ((<= targetDay) . date)