{-|
Module      : Cohort.Cohort
Description : Defines the Cohort and related types used in core logic of
              producing a Cohort from subject-level data.
Copyright   : (c) Target RWE 2023
License     : BSD3
Maintainer  : bbrown@targetrwe.com
              ljackman@targetrwe.com
              dpritchard@targetrwe.com
-}

{-# LANGUAGE DeriveGeneric #-}

module Cohort.Cohort where

import           Cohort.Criteria     (Criteria)
import           Cohort.IndexSet     (IndexSet)
import           Data.Aeson          (ToJSON, ToJSONKey)
import           Data.List           (sortOn)
import           Data.List.NonEmpty  (NonEmpty (..), (<|))
import           Data.Map.Strict     (Map, alter, empty, unionWith)
import           Data.Text           (Text)
import           EventDataTheory     (Event, Interval)
import           GHC.Generics        (Generic)
import           Variable

{- COHORT SPECIFICATION -}

-- | A container for all logic needed to construct a cohort. The provided
-- functions will be run on subject-level data, subject-by-subject. Logic will
-- be executed as part of a @'CohortApp'@ pipeline. The type 'CohortApp' itself
-- is not exported, but the user will create such a pipeline via the exported
-- 'cohortMain'.
--
-- Note 'runVariables' must construct *all* output variables. For a more
-- detailed explanation of the cohort-building pipeline, see the `Hasklepias`
-- module documentation from `hasklepias-main`.
data CohortSpec t m a
  = MkCohortSpec
      { forall t m a.
CohortSpec t m a -> NonEmpty (Event t m a) -> IndexSet a
runIndices   :: NonEmpty (Event t m a) -> IndexSet a
        -- ^ Constructs the set of index times, of type @Interval a@,
        -- for a given subject based on the full list of subject data in
        -- the 'NonEmpty' input list of @Event t m a@.
      , forall t m a.
CohortSpec t m a
-> NonEmpty (Event t m a) -> Interval a -> Criteria
runCriteria  :: NonEmpty (Event t m a) -> Interval a -> Criteria
        -- ^ Constructs a non-empty list of inclusion / exclusion
        -- criteria for a single subject, relative to a particular index
        -- time. There will be one 'Criteria' value for each subject and
        -- element of @'IndexSet' a@.
      , forall t m a.
CohortSpec t m a
-> NonEmpty (Event t m a) -> Interval a -> VariableRow
runVariables :: NonEmpty (Event t m a) -> Interval a -> VariableRow
        -- ^ Construct all output variables for the provided subject
        -- data, with respect to a particular index time.
      }

-- | Container for @CohortSpec@ for different cohorts, as given by the @Text@
-- keys.
type CohortSpecMap t m a = Map Text (CohortSpec t m a)

{- COHORT PIPELINE INPUT -}

-- | Wrapper for a @Text@ subject identifier. The wrapper allows us to
-- guarantee certain properties for subject identifiers, instead of allowing
-- arbitrary @Text@.
newtype SubjId
  = MkSubjId Text
  deriving (SubjId -> SubjId -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SubjId -> SubjId -> Bool
$c/= :: SubjId -> SubjId -> Bool
== :: SubjId -> SubjId -> Bool
$c== :: SubjId -> SubjId -> Bool
Eq, forall x. Rep SubjId x -> SubjId
forall x. SubjId -> Rep SubjId x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep SubjId x -> SubjId
$cfrom :: forall x. SubjId -> Rep SubjId x
Generic, Eq SubjId
SubjId -> SubjId -> Bool
SubjId -> SubjId -> Ordering
SubjId -> SubjId -> SubjId
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: SubjId -> SubjId -> SubjId
$cmin :: SubjId -> SubjId -> SubjId
max :: SubjId -> SubjId -> SubjId
$cmax :: SubjId -> SubjId -> SubjId
>= :: SubjId -> SubjId -> Bool
$c>= :: SubjId -> SubjId -> Bool
> :: SubjId -> SubjId -> Bool
$c> :: SubjId -> SubjId -> Bool
<= :: SubjId -> SubjId -> Bool
$c<= :: SubjId -> SubjId -> Bool
< :: SubjId -> SubjId -> Bool
$c< :: SubjId -> SubjId -> Bool
compare :: SubjId -> SubjId -> Ordering
$ccompare :: SubjId -> SubjId -> Ordering
Ord, Int -> SubjId -> ShowS
[SubjId] -> ShowS
SubjId -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SubjId] -> ShowS
$cshowList :: [SubjId] -> ShowS
show :: SubjId -> String
$cshow :: SubjId -> String
showsPrec :: Int -> SubjId -> ShowS
$cshowsPrec :: Int -> SubjId -> ShowS
Show)

instance ToJSON SubjId

-- | Subject-level data. The API via CohortSpec allows programmers to
-- manipulate contents of 'subjData' only, not 'subjId'. Attempting to process
-- a 'Subject' without any events is ill-defined, hence the use of 'NonEmpty'.
data Subject t m a
  = MkSubject
      { forall t m a. Subject t m a -> SubjId
subjId   :: SubjId
      , forall t m a. Subject t m a -> NonEmpty (Event t m a)
subjData :: NonEmpty (Event t m a)
      }
  deriving (Subject t m a -> Subject t m a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall t m a.
(Eq a, Eq t, Eq m) =>
Subject t m a -> Subject t m a -> Bool
/= :: Subject t m a -> Subject t m a -> Bool
$c/= :: forall t m a.
(Eq a, Eq t, Eq m) =>
Subject t m a -> Subject t m a -> Bool
== :: Subject t m a -> Subject t m a -> Bool
$c== :: forall t m a.
(Eq a, Eq t, Eq m) =>
Subject t m a -> Subject t m a -> Bool
Eq, Int -> Subject t m a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall t m a.
(Show t, Show m, Show a, Ord a) =>
Int -> Subject t m a -> ShowS
forall t m a.
(Show t, Show m, Show a, Ord a) =>
[Subject t m a] -> ShowS
forall t m a.
(Show t, Show m, Show a, Ord a) =>
Subject t m a -> String
showList :: [Subject t m a] -> ShowS
$cshowList :: forall t m a.
(Show t, Show m, Show a, Ord a) =>
[Subject t m a] -> ShowS
show :: Subject t m a -> String
$cshow :: forall t m a.
(Show t, Show m, Show a, Ord a) =>
Subject t m a -> String
showsPrec :: Int -> Subject t m a -> ShowS
$cshowsPrec :: forall t m a.
(Show t, Show m, Show a, Ord a) =>
Int -> Subject t m a -> ShowS
Show)

{- COHORT PIPELINE OUTPUT -}

-- | Internal. Status used for final inclusion/exclusion and the key of
-- 'attritionByStatus'. An observational unit can have multiple 'Exclude'
-- values but only be associated with a single 'ExcludeBY'.
data AttritionStatus
  = Included
  | ExcludedBy Text
  deriving (AttritionStatus -> AttritionStatus -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AttritionStatus -> AttritionStatus -> Bool
$c/= :: AttritionStatus -> AttritionStatus -> Bool
== :: AttritionStatus -> AttritionStatus -> Bool
$c== :: AttritionStatus -> AttritionStatus -> Bool
Eq, forall x. Rep AttritionStatus x -> AttritionStatus
forall x. AttritionStatus -> Rep AttritionStatus x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep AttritionStatus x -> AttritionStatus
$cfrom :: forall x. AttritionStatus -> Rep AttritionStatus x
Generic, Eq AttritionStatus
AttritionStatus -> AttritionStatus -> Bool
AttritionStatus -> AttritionStatus -> Ordering
AttritionStatus -> AttritionStatus -> AttritionStatus
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: AttritionStatus -> AttritionStatus -> AttritionStatus
$cmin :: AttritionStatus -> AttritionStatus -> AttritionStatus
max :: AttritionStatus -> AttritionStatus -> AttritionStatus
$cmax :: AttritionStatus -> AttritionStatus -> AttritionStatus
>= :: AttritionStatus -> AttritionStatus -> Bool
$c>= :: AttritionStatus -> AttritionStatus -> Bool
> :: AttritionStatus -> AttritionStatus -> Bool
$c> :: AttritionStatus -> AttritionStatus -> Bool
<= :: AttritionStatus -> AttritionStatus -> Bool
$c<= :: AttritionStatus -> AttritionStatus -> Bool
< :: AttritionStatus -> AttritionStatus -> Bool
$c< :: AttritionStatus -> AttritionStatus -> Bool
compare :: AttritionStatus -> AttritionStatus -> Ordering
$ccompare :: AttritionStatus -> AttritionStatus -> Ordering
Ord, Int -> AttritionStatus -> ShowS
[AttritionStatus] -> ShowS
AttritionStatus -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AttritionStatus] -> ShowS
$cshowList :: [AttritionStatus] -> ShowS
show :: AttritionStatus -> String
$cshow :: AttritionStatus -> String
showsPrec :: Int -> AttritionStatus -> ShowS
$cshowsPrec :: Int -> AttritionStatus -> ShowS
Show)

instance ToJSON AttritionStatus
instance ToJSONKey AttritionStatus

-- | Container for accumulated inclusion/exclusion counts.
data AttritionInfo
  = MkAttritionInfo
      { AttritionInfo -> Int
subjectsProcessed :: Int
      , AttritionInfo -> Int
unitsProcessed    :: Int
      , AttritionInfo -> Map AttritionStatus Int
attritionByStatus :: Map AttritionStatus Int
      }
  deriving (AttritionInfo -> AttritionInfo -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AttritionInfo -> AttritionInfo -> Bool
$c/= :: AttritionInfo -> AttritionInfo -> Bool
== :: AttritionInfo -> AttritionInfo -> Bool
$c== :: AttritionInfo -> AttritionInfo -> Bool
Eq, forall x. Rep AttritionInfo x -> AttritionInfo
forall x. AttritionInfo -> Rep AttritionInfo x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep AttritionInfo x -> AttritionInfo
$cfrom :: forall x. AttritionInfo -> Rep AttritionInfo x
Generic, Int -> AttritionInfo -> ShowS
[AttritionInfo] -> ShowS
AttritionInfo -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AttritionInfo] -> ShowS
$cshowList :: [AttritionInfo] -> ShowS
show :: AttritionInfo -> String
$cshow :: AttritionInfo -> String
showsPrec :: Int -> AttritionInfo -> ShowS
$cshowsPrec :: Int -> AttritionInfo -> ShowS
Show)

instance ToJSON AttritionInfo

-- | Identifier for the unit element of a Cohort, @ObsUnit@. @fromSubjId@
-- should refer to the @subjId@ of the @Subject@ from which an @ObsUnit@ was
-- constructed. There should be one @ObsId@, and hence one @ObsUnit@, per pair
-- formed from a given @Subject@ and each element of the subject's @IndexSet@.
-- @indexTime@ gives the element of @IndexSet@ associated with this @ObsUnit@.
data ObsId a
  = MkObsId
      { forall a. ObsId a -> SubjId
fromSubjId :: SubjId
      , forall a. ObsId a -> Interval a
indexTime  :: Interval a
      }
  deriving (ObsId a -> ObsId a -> Bool
forall a. Eq a => ObsId a -> ObsId a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ObsId a -> ObsId a -> Bool
$c/= :: forall a. Eq a => ObsId a -> ObsId a -> Bool
== :: ObsId a -> ObsId a -> Bool
$c== :: forall a. Eq a => ObsId a -> ObsId a -> Bool
Eq, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (ObsId a) x -> ObsId a
forall a x. ObsId a -> Rep (ObsId a) x
$cto :: forall a x. Rep (ObsId a) x -> ObsId a
$cfrom :: forall a x. ObsId a -> Rep (ObsId a) x
Generic, ObsId a -> ObsId a -> Bool
ObsId a -> ObsId a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (ObsId a)
forall a. Ord a => ObsId a -> ObsId a -> Bool
forall a. Ord a => ObsId a -> ObsId a -> Ordering
forall a. Ord a => ObsId a -> ObsId a -> ObsId a
min :: ObsId a -> ObsId a -> ObsId a
$cmin :: forall a. Ord a => ObsId a -> ObsId a -> ObsId a
max :: ObsId a -> ObsId a -> ObsId a
$cmax :: forall a. Ord a => ObsId a -> ObsId a -> ObsId a
>= :: ObsId a -> ObsId a -> Bool
$c>= :: forall a. Ord a => ObsId a -> ObsId a -> Bool
> :: ObsId a -> ObsId a -> Bool
$c> :: forall a. Ord a => ObsId a -> ObsId a -> Bool
<= :: ObsId a -> ObsId a -> Bool
$c<= :: forall a. Ord a => ObsId a -> ObsId a -> Bool
< :: ObsId a -> ObsId a -> Bool
$c< :: forall a. Ord a => ObsId a -> ObsId a -> Bool
compare :: ObsId a -> ObsId a -> Ordering
$ccompare :: forall a. Ord a => ObsId a -> ObsId a -> Ordering
Ord, Int -> ObsId a -> ShowS
forall a. (Show a, Ord a) => Int -> ObsId a -> ShowS
forall a. (Show a, Ord a) => [ObsId a] -> ShowS
forall a. (Show a, Ord a) => ObsId a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ObsId a] -> ShowS
$cshowList :: forall a. (Show a, Ord a) => [ObsId a] -> ShowS
show :: ObsId a -> String
$cshow :: forall a. (Show a, Ord a) => ObsId a -> String
showsPrec :: Int -> ObsId a -> ShowS
$cshowsPrec :: forall a. (Show a, Ord a) => Int -> ObsId a -> ShowS
Show)

instance (ToJSON a) => ToJSON (ObsId a)

-- | There should be one @ObsUnit@ per pair formed from a given @Subject@ and
-- each element of the subject's @IndexSet@.
data ObsUnit a
  = MkObsUnit
      { forall a. ObsUnit a -> ObsId a
obsId   :: ObsId a
      , forall a. ObsUnit a -> VariableRow
obsData :: VariableRow
      }
  deriving (Int -> ObsUnit a -> ShowS
forall a. (Show a, Ord a) => Int -> ObsUnit a -> ShowS
forall a. (Show a, Ord a) => [ObsUnit a] -> ShowS
forall a. (Show a, Ord a) => ObsUnit a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ObsUnit a] -> ShowS
$cshowList :: forall a. (Show a, Ord a) => [ObsUnit a] -> ShowS
show :: ObsUnit a -> String
$cshow :: forall a. (Show a, Ord a) => ObsUnit a -> String
showsPrec :: Int -> ObsUnit a -> ShowS
$cshowsPrec :: forall a. (Show a, Ord a) => Int -> ObsUnit a -> ShowS
Show)

-- TODO it should be an error that a subject does not have an index date. such
-- cases were dropped, both previously and in the current implementation. an
-- error-handling redesign should address that.

-- | Internal type to hold the result of @Cohort.Core.'evaluateSubj'@. It
-- should not be used for any other purpose.
data EvaluatedSubject a
  = SNoIndex SubjId
  -- ^ @Subject@ had no index. These contribute to the subjectsProcessed count
  -- but are silently dropped.
  | SUnits (SUnitData a)
  deriving (Int -> EvaluatedSubject a -> ShowS
forall a. (Show a, Ord a) => Int -> EvaluatedSubject a -> ShowS
forall a. (Show a, Ord a) => [EvaluatedSubject a] -> ShowS
forall a. (Show a, Ord a) => EvaluatedSubject a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EvaluatedSubject a] -> ShowS
$cshowList :: forall a. (Show a, Ord a) => [EvaluatedSubject a] -> ShowS
show :: EvaluatedSubject a -> String
$cshow :: forall a. (Show a, Ord a) => EvaluatedSubject a -> String
showsPrec :: Int -> EvaluatedSubject a -> ShowS
$cshowsPrec :: forall a. (Show a, Ord a) => Int -> EvaluatedSubject a -> ShowS
Show)

-- | Internal. Holds data for 'SUnits'.
data SUnitData a
  = MkSUnitData
      { forall a. SUnitData a -> [(ObsId a, Text)]
excludeUnits  :: [(ObsId a, Text)]
        -- ^ Observational units that are to be excluded from the cohort.
        -- Only @ObsId@ is retained since no @obsData@ is to be computed.
        -- The @Text@ should be the @statusLabel@ field's value of
        -- @Criterion@.
      , forall a. SUnitData a -> [ObsUnit a]
includeUnits  :: [ObsUnit a]
        -- ^ Observational units to be included in the cohort, on which
        -- @runVariables@ is to be computed.
      , forall a. SUnitData a -> AttritionInfo
attritionSubj :: AttritionInfo
        -- ^ Subject-level attrition, to allow computing attritionInfo as each
        -- subject is processed.
      }
  deriving (Int -> SUnitData a -> ShowS
forall a. (Show a, Ord a) => Int -> SUnitData a -> ShowS
forall a. (Show a, Ord a) => [SUnitData a] -> ShowS
forall a. (Show a, Ord a) => SUnitData a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SUnitData a] -> ShowS
$cshowList :: forall a. (Show a, Ord a) => [SUnitData a] -> ShowS
show :: SUnitData a -> String
$cshow :: forall a. (Show a, Ord a) => SUnitData a -> String
showsPrec :: Int -> SUnitData a -> ShowS
$cshowsPrec :: forall a. (Show a, Ord a) => Int -> SUnitData a -> ShowS
Show)

-- | Accumulated data from a single list of subjects, processed according to
-- the logic provided in a @CohortSpec@. It is the output type of @evalCohort@
-- and should be produced nowhere else.
data Cohort a
  = MkCohort
      { forall a. Cohort a -> AttritionInfo
attritionInfo :: AttritionInfo
      , forall a. Cohort a -> [ObsUnit a]
cohortData    :: [ObsUnit a]
      }
  deriving (Int -> Cohort a -> ShowS
forall a. (Show a, Ord a) => Int -> Cohort a -> ShowS
forall a. (Show a, Ord a) => [Cohort a] -> ShowS
forall a. (Show a, Ord a) => Cohort a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Cohort a] -> ShowS
$cshowList :: forall a. (Show a, Ord a) => [Cohort a] -> ShowS
show :: Cohort a -> String
$cshow :: forall a. (Show a, Ord a) => Cohort a -> String
showsPrec :: Int -> Cohort a -> ShowS
$cshowsPrec :: forall a. (Show a, Ord a) => Int -> Cohort a -> ShowS
Show)

type CohortMap a = Map Text (Cohort a)

{- EXPORTED UTILITIES -}

-- NOTE this does not use Data.Map.fromListWith or similar because it should be
-- similar in complexity but avoids the (unlikely) issue related to map size
-- mentioned in that module. Should be evalated in an optimization pass.

-- TODO as part of a future redesign providing error-handling, this function
-- should incorporate some minimal validation of subject ids, which possibly is
-- configurable.

-- | Accumulate @[(Text, Event t m a)]@ into @[Subject t m a]@. The former is
-- the \'flat\' output format produced in @EventDataTheory.parseEventLinesL'@.
eventsToSubject :: [(Text, Event t m a)] -> [Subject t m a]
eventsToSubject :: forall t m a. [(Text, Event t m a)] -> [Subject t m a]
eventsToSubject [] = []
eventsToSubject [(Text, Event t m a)]
es = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {t} {m} {a}.
(Text, Event t m a) -> [Subject t m a] -> [Subject t m a]
op [] [(Text, Event t m a)]
es'
  where es' :: [(Text, Event t m a)]
es' = forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Text, Event t m a)]
es
        op :: (Text, Event t m a) -> [Subject t m a] -> [Subject t m a]
op (Text
sid', Event t m a
e') [] = [forall t m a. SubjId -> NonEmpty (Event t m a) -> Subject t m a
MkSubject (Text -> SubjId
MkSubjId Text
sid') (Event t m a
e' forall a. a -> [a] -> NonEmpty a
:| [])]
        op (Text
sid', Event t m a
e') (subj :: Subject t m a
subj@(MkSubject (MkSubjId Text
sid) NonEmpty (Event t m a)
ess) : [Subject t m a]
ss)
          | Text
sid forall a. Eq a => a -> a -> Bool
== Text
sid' = forall t m a. SubjId -> NonEmpty (Event t m a) -> Subject t m a
MkSubject (Text -> SubjId
MkSubjId Text
sid) (Event t m a
e' forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty (Event t m a)
ess) forall a. a -> [a] -> [a]
: [Subject t m a]
ss
          | Bool
otherwise = forall t m a. SubjId -> NonEmpty (Event t m a) -> Subject t m a
MkSubject (Text -> SubjId
MkSubjId Text
sid') (Event t m a
e' forall a. a -> [a] -> NonEmpty a
:| []) forall a. a -> [a] -> [a]
: Subject t m a
subj forall a. a -> [a] -> [a]
: [Subject t m a]
ss


{- INTERNAL UTILITIES -}

emptyCohort :: Cohort b
emptyCohort :: forall b. Cohort b
emptyCohort = forall a. AttritionInfo -> [ObsUnit a] -> Cohort a
MkCohort AttritionInfo
emptyAttrition []

emptyAttrition :: AttritionInfo
emptyAttrition :: AttritionInfo
emptyAttrition = Int -> Int -> Map AttritionStatus Int -> AttritionInfo
MkAttritionInfo Int
0 Int
0 forall k a. Map k a
empty

alterAttrition :: AttritionStatus -> Map AttritionStatus Int -> Map AttritionStatus Int
alterAttrition :: AttritionStatus
-> Map AttritionStatus Int -> Map AttritionStatus Int
alterAttrition = forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
alter forall {a}. Num a => Maybe a -> Maybe a
op
  where
        -- If key `k` doesn't exist, insert it with value 1.
        -- Else increment by 1.
        op :: Maybe a -> Maybe a
op Maybe a
Nothing  = forall a. a -> Maybe a
Just a
1
        op (Just a
n) = forall a. a -> Maybe a
Just (a
nforall a. Num a => a -> a -> a
+a
1)

-- | Internal. Increment AttritionInfo with a single unit's information. Does
-- not alter @subjectsProcessed@.
incrementAttritionForUnit :: AttritionStatus -> AttritionInfo -> AttritionInfo
incrementAttritionForUnit :: AttritionStatus -> AttritionInfo -> AttritionInfo
incrementAttritionForUnit AttritionStatus
s AttritionInfo
info = AttritionInfo
info{ unitsProcessed :: Int
unitsProcessed = Int
1 forall a. Num a => a -> a -> a
+ AttritionInfo -> Int
unitsProcessed AttritionInfo
info, attritionByStatus :: Map AttritionStatus Int
attritionByStatus = AttritionStatus
-> Map AttritionStatus Int -> Map AttritionStatus Int
alterAttrition AttritionStatus
s (AttritionInfo -> Map AttritionStatus Int
attritionByStatus AttritionInfo
info) }

-- | Combine two @AttritionInfo@ values by summing the counts, and taking the
-- union of @attritionByStatus@, summing values for shared keys.
combineAttrition :: AttritionInfo -> AttritionInfo -> AttritionInfo
combineAttrition :: AttritionInfo -> AttritionInfo -> AttritionInfo
combineAttrition AttritionInfo
a1 AttritionInfo
a2 = MkAttritionInfo { subjectsProcessed :: Int
subjectsProcessed = Int
sp, unitsProcessed :: Int
unitsProcessed = Int
up, attritionByStatus :: Map AttritionStatus Int
attritionByStatus = Map AttritionStatus Int
m }
  where
    sp :: Int
sp = AttritionInfo -> Int
subjectsProcessed AttritionInfo
a1 forall a. Num a => a -> a -> a
+ AttritionInfo -> Int
subjectsProcessed AttritionInfo
a2
    up :: Int
up = AttritionInfo -> Int
unitsProcessed AttritionInfo
a1 forall a. Num a => a -> a -> a
+ AttritionInfo -> Int
unitsProcessed AttritionInfo
a2
    m :: Map AttritionStatus Int
m = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWith forall a. Num a => a -> a -> a
(+) (AttritionInfo -> Map AttritionStatus Int
attritionByStatus AttritionInfo
a1) (AttritionInfo -> Map AttritionStatus Int
attritionByStatus AttritionInfo
a2)