{-|
Module      : Cohort.Core
Description : Internal module defining logic to process a list of @Subject@
              into a @Cohort@.
Copyright   : (c) Target RWE 2023
License     : BSD3
Maintainer  : bbrown@targetrwe.com 
              ljackman@targetrwe.com 
              dpritchard@targetrwe.com

This is module supplies the logic for how cohorts are processed based on the
user-defined @CohortSpecMap@. It is intended to be used internally to
`hasklepias-main` and not re-exported from there.

For high-level documentation of the cohort-processing pipeline, see the
`Hasklepias` module documentation from `hasklepias-main`.

Command-line options, defined in `hasklepias-main`, should not affect the
internal logic of cohort processing. Only the CohortSpec should be able to do so.

Note 'subjData' is processed *only* via the logic provided in 'CohortSpec'.
-}

module Cohort.Core where

import           Cohort.Cohort
import           Cohort.Criteria    (Criterion (..), firstExclude)
import qualified Cohort.IndexSet    as IS
import           Data.Foldable      (foldl')
import qualified Data.Map.Strict    as M
import           Data.Text          (Text)
import           EventDataTheory    (Interval)

{- Subject-level evaluation -}

-- | Evaluate a subject at a single index time, producing @Left@ if there were
-- any @Exclude@ among the computed @Criteria@, and otherwise returning @Right
-- (ObsUnit b)@ with variables computed via @runVariables@.
evalSubjAtIndex :: CohortSpec t m a -> Subject t m a -> Interval a -> Either (ObsId a, Text) (ObsUnit a)
evalSubjAtIndex :: forall t m a.
CohortSpec t m a
-> Subject t m a
-> Interval a
-> Either (ObsId a, Text) (ObsUnit a)
evalSubjAtIndex CohortSpec t m a
spec Subject t m a
subj Interval a
i = forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall a b. b -> Either a b
Right ObsUnit a
unit) (\Criterion
c -> forall a b. a -> Either a b
Left (ObsId a
obsid, Criterion -> Text
statusLabel Criterion
c)) forall a b. (a -> b) -> a -> b
$ Criteria -> Maybe Criterion
firstExclude Criteria
crits
  where
    sid :: SubjId
sid = forall t m a. Subject t m a -> SubjId
subjId Subject t m a
subj
    dt :: NonEmpty (Event t m a)
dt = forall t m a. Subject t m a -> NonEmpty (Event t m a)
subjData Subject t m a
subj
    obsid :: ObsId a
obsid = forall a. SubjId -> Interval a -> ObsId a
MkObsId SubjId
sid Interval a
i
    crits :: Criteria
crits = forall t m a.
CohortSpec t m a
-> NonEmpty (Event t m a) -> Interval a -> Criteria
runCriteria CohortSpec t m a
spec NonEmpty (Event t m a)
dt Interval a
i
    -- NOTE this step is expensive if runVariables is expensive
    unit :: ObsUnit a
unit = forall a. ObsId a -> VariableRow -> ObsUnit a
MkObsUnit ObsId a
obsid forall a b. (a -> b) -> a -> b
$ forall t m a.
CohortSpec t m a
-> NonEmpty (Event t m a) -> Interval a -> VariableRow
runVariables CohortSpec t m a
spec NonEmpty (Event t m a)
dt Interval a
i

-- | Compute @IndexSet@, @Criteria@ and variables for a given subject via
-- strict fold on the @IndexSet@. Returns @SNoIndex@ if and only if @IS.'null'
-- idxs@. As reflected in the definition of @EvaluatedSubject@, @runVariables@
-- is computed only for units whose @Criteria@ have @status@ @Include@. This
-- also updates the @attritionSubj@, the @AttritionInfo@ associated with this
-- subject.
evalSubj :: CohortSpec t m a -> Subject t m a -> EvaluatedSubject a
evalSubj :: forall t m a.
CohortSpec t m a -> Subject t m a -> EvaluatedSubject a
evalSubj CohortSpec t m a
spec Subject t m a
subj = forall a b. (a -> Interval b -> a) -> a -> IndexSet b -> a
IS.foldl' EvaluatedSubject a -> Interval a -> EvaluatedSubject a
op (forall a. SubjId -> EvaluatedSubject a
SNoIndex (forall t m a. Subject t m a -> SubjId
subjId Subject t m a
subj)) IndexSet a
idxs
  where
        idxs :: IndexSet a
idxs = forall t m a.
CohortSpec t m a -> NonEmpty (Event t m a) -> IndexSet a
runIndices CohortSpec t m a
spec (forall t m a. Subject t m a -> NonEmpty (Event t m a)
subjData Subject t m a
subj)
        -- Single-subject attrition. Accumulates units processed and attrition
        -- by status.
        initAttrition :: AttritionInfo
initAttrition = AttritionInfo
emptyAttrition{ subjectsProcessed :: Int
subjectsProcessed = Int
1 }
        updateExcl :: (ObsId a, Text)
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateExcl (ObsId a
oid, Text
t) [(ObsId a, Text)]
ess [ObsUnit a]
iss AttritionInfo
attr = forall a. SUnitData a -> EvaluatedSubject a
SUnits forall a b. (a -> b) -> a -> b
$ forall a.
[(ObsId a, Text)] -> [ObsUnit a] -> AttritionInfo -> SUnitData a
MkSUnitData ((ObsId a
oid, Text
t)forall a. a -> [a] -> [a]
:[(ObsId a, Text)]
ess) [ObsUnit a]
iss forall a b. (a -> b) -> a -> b
$
                                      AttritionStatus -> AttritionInfo -> AttritionInfo
incrementAttritionForUnit (Text -> AttritionStatus
ExcludedBy Text
t) AttritionInfo
attr
        updateIncl :: ObsUnit a
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateIncl ObsUnit a
x [(ObsId a, Text)]
ess [ObsUnit a]
iss AttritionInfo
attr = forall a. SUnitData a -> EvaluatedSubject a
SUnits forall a b. (a -> b) -> a -> b
$ forall a.
[(ObsId a, Text)] -> [ObsUnit a] -> AttritionInfo -> SUnitData a
MkSUnitData [(ObsId a, Text)]
ess (ObsUnit a
xforall a. a -> [a] -> [a]
:[ObsUnit a]
iss) forall a b. (a -> b) -> a -> b
$
                                      AttritionStatus -> AttritionInfo -> AttritionInfo
incrementAttritionForUnit AttritionStatus
Included AttritionInfo
attr
        op :: EvaluatedSubject a -> Interval a -> EvaluatedSubject a
op (SNoIndex SubjId
_) Interval a
i = case forall t m a.
CohortSpec t m a
-> Subject t m a
-> Interval a
-> Either (ObsId a, Text) (ObsUnit a)
evalSubjAtIndex CohortSpec t m a
spec Subject t m a
subj Interval a
i of
                              Left (ObsId a, Text)
x  -> forall {a}.
(ObsId a, Text)
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateExcl (ObsId a, Text)
x [] [] AttritionInfo
initAttrition
                              Right ObsUnit a
x -> forall {a}.
ObsUnit a
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateIncl ObsUnit a
x [] [] AttritionInfo
initAttrition
        op (SUnits (MkSUnitData [(ObsId a, Text)]
ess [ObsUnit a]
iss AttritionInfo
attr)) Interval a
i = case forall t m a.
CohortSpec t m a
-> Subject t m a
-> Interval a
-> Either (ObsId a, Text) (ObsUnit a)
evalSubjAtIndex CohortSpec t m a
spec Subject t m a
subj Interval a
i of
                                  Left (ObsId a, Text)
x  -> forall {a}.
(ObsId a, Text)
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateExcl (ObsId a, Text)
x [(ObsId a, Text)]
ess [ObsUnit a]
iss AttritionInfo
attr
                                  Right ObsUnit a
x -> forall {a}.
ObsUnit a
-> [(ObsId a, Text)]
-> [ObsUnit a]
-> AttritionInfo
-> EvaluatedSubject a
updateIncl ObsUnit a
x [(ObsId a, Text)]
ess [ObsUnit a]
iss AttritionInfo
attr


-- | Processes the cohort-building logic provided in @CohortSpec@ on a list of
-- @Subject@s. @Subject@s are processes one-by-one, and @AttritionInfo@ is
-- accumulated as each subject is processed. @Subject@s whose set of index
-- times is empty are counted only in the @subjectsProcessed@ field of
-- @AttritionInfo@.
--
-- Each @Subject@ produces one @ObsUnit@ in the @Cohort@ for each element of
-- the subject's @IndexSet@ such that all @Criteria@ associated with that
-- index time have status @Include@. If at least one element of the @Criteria@
-- for an index time has status @Exclude@, no variables are computed and the
-- exclusion is recorded in the @AttritionInfo@.
evalCohort :: CohortSpec t m a -> [Subject t m a] -> Cohort a
evalCohort :: forall t m a. CohortSpec t m a -> [Subject t m a] -> Cohort a
evalCohort CohortSpec t m a
spec [Subject t m a]
ss = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {a}. Cohort a -> EvaluatedSubject a -> Cohort a
op forall b. Cohort b
emptyCohort (forall a b. (a -> b) -> [a] -> [b]
map (forall t m a.
CohortSpec t m a -> Subject t m a -> EvaluatedSubject a
evalSubj CohortSpec t m a
spec) [Subject t m a]
ss)
  where
    -- SNoIndex is counted for subjects processed but nothing more.
    emptySubjAttrition :: AttritionInfo
emptySubjAttrition = AttritionInfo
emptyAttrition{ subjectsProcessed :: Int
subjectsProcessed = Int
1 }
    op :: Cohort a -> EvaluatedSubject a -> Cohort a
op (MkCohort AttritionInfo
a [ObsUnit a]
us) (SNoIndex SubjId
_) = forall a. AttritionInfo -> [ObsUnit a] -> Cohort a
MkCohort (AttritionInfo -> AttritionInfo -> AttritionInfo
combineAttrition AttritionInfo
a AttritionInfo
emptySubjAttrition) [ObsUnit a]
us
    -- NOTE: this concat is a good place to look for efficiency gains if needed.
    -- us1 ++ us2 preserves order of subjects.
    op (MkCohort AttritionInfo
a1 [ObsUnit a]
us1) (SUnits (MkSUnitData [(ObsId a, Text)]
_ [ObsUnit a]
us2 AttritionInfo
a2)) = forall a. AttritionInfo -> [ObsUnit a] -> Cohort a
MkCohort (AttritionInfo -> AttritionInfo -> AttritionInfo
combineAttrition AttritionInfo
a1 AttritionInfo
a2) ([ObsUnit a]
us1 forall a. [a] -> [a] -> [a]
++ [ObsUnit a]
us2)


-- | A convenience function to apply @evalCohort@ to a fixed list of subjects,
-- once for each element of a @CohortMapSpec@. Each element of the
-- @CohortMapSpec@ represents a different, user-provided cohort-building
-- logic.
evalCohortMap :: CohortSpecMap t m a -> [Subject t m a] -> CohortMap a
evalCohortMap :: forall t m a. CohortSpecMap t m a -> [Subject t m a] -> CohortMap a
evalCohortMap CohortSpecMap t m a
specm [Subject t m a]
subj = forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall t m a. CohortSpec t m a -> [Subject t m a] -> Cohort a
`evalCohort` [Subject t m a]
subj) CohortSpecMap t m a
specm