interval-algebra foundations

import IntervalAlgebra
import IntervalAlgebra.IntervalDiagram

import Data.Bifunctor ( Bifunctor(..) )
import Data.List ( sort )
import Data.Time
    ( addDays, fromGregorian, secondsToDiffTime, Day, UTCTime(..) )
import Witch ( into )
import Data.Set ( fromList, difference, Set )

data Interval a

-- An example Interval Integer
print ivInteger
---> (2, 6)

-- An example Interval Day
print ivDay
---> (1967-01-18, 1967-01-24)

-- An example Interval UTCTime
print ivUTC
---> (1967-01-18 09:00:00 UTC, 1967-01-18 09:13:20 UTC)

Eq a => Eq (Interval a)
Ord a => Ord (Interval a)
instance (Show a, Ord a) => Show (Interval a)

print $ ivInteger == ivInteger
---> True

print $ ivDay < ivDay
---> False

print $ show ivInteger
---> "(2, 6)"

iv0to2, iv2to4, iv2to5, iv4to5, iv5to8, iv6to8, iv3to6 :: Interval Integer
iv0to2 = safeInterval (0, 2)
iv2to4 = safeInterval (2, 4)
iv2to5 = safeInterval (2, 5)
iv3to6 = safeInterval (3, 6)
iv4to5 = safeInterval (4, 5)
iv5to8 = safeInterval (5, 8)
iv6to8 = safeInterval (6, 8)

data PairedInterval b a

makePairedInterval :: b -> Interval a -> PairedInterval b a

pairListstringInteger :: PairedInterval [String] Integer
pairListstringInteger =
  makePairedInterval ["John", "Paul", "George", "Ringo"] ivInteger

pairStringDay :: PairedInterval String Day
pairStringDay = makePairedInterval "vacation" ivDay

print pairListstringInteger
---> {(2, 6), ["John","Paul","George","Ringo"]}

print pairStringDay
---> {(1967-01-18, 1967-01-24), "vacation"}

(Eq a, Eq b) => Eq (PairedInterval b a)
(Eq a, Eq b, Ord a) => Ord (PairedInterval b a)
(Show b, Show a, Ord a) => Show (PairedInterval b a)

print $ pairStringDay == pairStringDay
---> True

print $ pairListstringInteger < pairListstringInteger
---> False

print $ show pairStringDay
---> "{(1967-01-18, 1967-01-24), \"vacation\"}"

class Intervallic i where

getInterval :: i a -> Interval a
setInterval :: i a -> Interval b -> i b

begin :: Intervallic i => i a -> a
end :: Intervallic i => i a -> a

class (Ord a, Num b, Ord b) => IntervalSizeable a b | a -> b

moment :: forall a . b
duration :: Intervallic i => i a -> b
add :: b -> a -> a
diff :: a -> a -> b

class Intervallic i => IntervalCombinable i a

(.+.) :: i a -> i a -> Maybe (i a)
(><) :: i a -> i a -> Maybe (i a)
(<+>) :: (Semigroup (f (i a)), Applicative f) => i a -> i a -> f (i a)

print ivInteger
---> (2, 6)

print $ expandl 4 ivInteger
---> (-2, 6)

print $ expandl 0 ivInteger
---> (2, 6)

print $ expandr 5 ivInteger
---> (2, 11)

print $ expandr (-3) ivInteger
---> (2, 6)

print $ expand 4 5 ivInteger
---> (-2, 11)

print $ expand 0 (-3) ivInteger
---> (2, 6)

print ivInteger
---> (2, 6)

beginervalFromEnd 5 ivInteger
---> (6, 11)

beginervalFromEnd (-2) ivInteger
---> (6, 7)

endervalFromBegin 12 ivInteger
---> (-10, 2)

endervalFromBegin (-6) ivInteger
---> (1, 2)

print $ momentize ivInteger
---> (2, 3)

shiftFromBegin
  :: (IntervalSizeable a b, Intervallic i0 a)
  => i0 a -> i1 a -> i1 b

shiftFromEnd
  :: (IntervalSizeable a b, Intervallic i0 a)
  => i0 a -> i1 a -> i1 b

print [iv2to4, iv5to8]
---> [(2, 4),(5, 8)]

print ivDay
---> (1967-01-18, 1967-01-24)

print $ shiftFromBegin iv2to4 iv5to8
---> (3, 6)

print $ shiftFromBegin ivDay ivDay
---> (0, 6)

print $ shiftFromEnd iv2to4 iv5to8
---> (1, 4)

print $ shiftFromEnd ivDay ivDay
---> (-6, 0)

type ComparativePredicateOf1 a = (a -> a -> Bool)
type ComparativePredicateOf2 a b = (a -> b -> Bool)

data IntervalRelation =
    Before        -- before/precedes
  | Meets         -- meets
  | Overlaps      -- overlaps
  | FinishedBy    -- finishedBy
  | Contains      -- contains
  | Starts        -- starts
  | Equals        -- equals
  | StartedBy     -- startedBy
  | During        -- during
  | Finishes      -- finishes
  | OverlappedBy  -- overlappedBy
  | MetBy         -- metBy
  | After         -- after/precededBy

--           <- [iv0to2]
  --         <- [iv2to4]
  ---        <- [iv2to5]
   ---       <- [iv3to6]
    -        <- [iv4to5]
      --     <- [iv6to8]
     ---     <- [iv5to8]
============

(Ord a, Intervallic i0, Intervallic i1) => ComparativePredicateOf2 (i0 a) (i1 a)

data DataType = Enrollment | Treatment TreatmentType | Diagnosis DiagnosisType
  deriving (Eq, Ord, Show)

data TreatmentType = StandardVaccine | NewVaccine
  deriving (Eq, Ord, Show)

data DiagnosisType = RightAsRain | UpsetTummy | CommonCold | Flu
  deriving (Eq, Ord, Show)

type StudyEvent = PairedInterval DataType Integer

type SubjEvents = [StudyEvent]

data ProcessedSubj = ProcessedSubj
  { getEnrollment   :: Maybe (Interval Integer)
  , getFirstTrt     :: Maybe StudyEvent
  , getFirstFlu     :: Maybe StudyEvent
  , getTrtType      :: Maybe TreatmentType
  , getTimeToFlu    :: Maybe Integer
  , getTimeToEndEnr :: Maybe Integer
  }
  deriving Show

maxEnrGap :: Integer
maxEnrGap = 8

-- Construct the elements of a `ProcessedSubj` one step at-a-time. Most of the
-- actual work is done by the helper functions defined below
processSubj :: SubjEvents -> ProcessedSubj
processSubj xs =
  let enrPeriod    = calcEnrPeriod xs                   -- enrollment period
      enrEvents    = calcEnrEvents enrPeriod xs         -- events within enrollment
      firstTrt     = findFirstTrt enrPeriod xs          -- first treatment in enr
      firstTrtType = extractTrtType firstTrt            -- first trt type
      firstTrtIv   = fmap getInterval firstTrt          -- first trt interval
      firstFlu     = findFirstFlu firstTrtIv enrEvents  -- first flu in enr
      ttFlu        = calcDiff firstFlu firstTrt         -- time to first flu
      ttEndEnr     = calcAtRisk enrPeriod firstTrt      -- time to end of enr
  in  ProcessedSubj enrPeriod firstTrt firstFlu firstTrtType ttFlu ttEndEnr

-- Construct the "enrollment period", which is defined at the period of time
-- with the start endpoint given by their earliest enrollment period and end
-- endpoint given by the first time that they fall out of the grace period. In
-- the event that a subject did not have any enrollment periods then the return
-- value is a Nothing.
--
-- Note that this function uses the `combineIntervals` function which was not
-- covered in this tutorial, but is exported from IntervalAlgebra via
-- IntervalAlgebra.IntervalUtilities.
calcEnrPeriod :: SubjEvents -> Maybe (Interval Integer)
calcEnrPeriod xs | null combinedPeriods = Nothing
                 | otherwise            = Just (head combinedPeriods)
  where combinedPeriods = (combineIntervals . addMaxEnrGap . extractEnrIvs) xs

-- Filter the enrollment events in the SubjEvents and extract the Interval
-- from each one
extractEnrIvs :: SubjEvents -> [Interval Integer]
extractEnrIvs = intervals . filter (checkEnr . getPairData)
 where
  checkEnr Enrollment = True
  checkEnr _          = False

-- Extend the end endpoint for each Interval by `maxEnrGap`
addMaxEnrGap :: [Interval Integer] -> [Interval Integer]
addMaxEnrGap = map (expandr maxEnrGap)

-- Filter the SubjEvents to those with endpoints that do not extend past the
-- enrollment period's endpoints
calcEnrEvents :: Maybe (Interval Integer) -> SubjEvents -> SubjEvents
calcEnrEvents mayIv xs = case mayIv of
  Nothing -> []
  Just y  -> filter (\x -> getInterval x `enclosedBy` y) xs

-- Find the first flu vaccine administrations occuring within the enrollment
-- period
findFirstTrt :: Maybe (Interval Integer) -> SubjEvents -> Maybe StudyEvent
findFirstTrt Nothing _ = Nothing
findFirstTrt (Just iv) xs | null filteredIntakes = Nothing
                          | otherwise            = Just (head filteredIntakes)
 where
  p x = checkTrt (getPairData x) && (getInterval x `enclosedBy` iv)
  checkTrt (Treatment _) = True
  checkTrt _             = False
  filteredIntakes = filter p xs

-- Extract the treatment type out of a StudyEvent. If there is no event or the
-- StudyEvent type isn't TreatmentType then return Nothing
extractTrtType :: Maybe StudyEvent -> Maybe TreatmentType
extractTrtType Nothing  = Nothing
extractTrtType (Just x) = case getPairData x of
  Treatment StandardVaccine -> Just StandardVaccine
  Treatment NewVaccine      -> Just NewVaccine
  _                         -> Nothing

-- Find the first flu diagnosis occuring within the enrollment period
findFirstFlu :: Maybe (Interval Integer) -> SubjEvents -> Maybe StudyEvent
findFirstFlu Nothing _ = Nothing
findFirstFlu (Just iv) xs | null filteredFlus = Nothing
                          | otherwise         = Just (head filteredFlus)
 where
  endedPrior = before <|> meets
  p x = checkDiagFlu (getPairData x) && (iv `endedPrior` getInterval x)
  checkDiagFlu (Diagnosis Flu) = True
  checkDiagFlu _               = False
  filteredFlus = filter p xs

-- Calculate the difference between the start endpoint of the first Intervallic
-- and the start endpoint of the second Intervallic
calcDiff
  :: (IntervalSizeable a b, Intervallic i0, Intervallic i1)
  => Maybe (i0 a)
  -> Maybe (i1 a)
  -> Maybe b
calcDiff (Just y) (Just x) = Just $ diff (begin y) (end x)
calcDiff _ _               = Nothing

-- Calculate the difference between the end endpoint of the first Intervallic
-- and the start endpoint of the second Intervallic
calcAtRisk
  :: (IntervalSizeable a b, Intervallic i0, Intervallic i1)
  => Maybe (i0 a)
  -> Maybe (i1 a)
  -> Maybe b
calcAtRisk (Just y) (Just x) = Just $ diff (end y) (end x)
calcAtRisk _ _               = Nothing

results :: [ProcessedSubj]
results = map processSubj [id1Events, id2Events, id3Events, id4Events]

print $ head results
---> ProcessedSubj {getEnrollment = Just (6, 430), getFirstTrt = Just {(22, 23), Treatment NewVaccine}, getFirstFlu = Nothing, getTrtType = Just NewVaccine, getTimeToFlu = Nothing, getTimeToEndEnr = Just 407}

print $ results !! 1
---> ProcessedSubj {getEnrollment = Just (2, 214), getFirstTrt = Just {(98, 99), Treatment StandardVaccine}, getFirstFlu = Nothing, getTrtType = Just StandardVaccine, getTimeToFlu = Nothing, getTimeToEndEnr = Just 115}

print $ results !! 2
---> ProcessedSubj {getEnrollment = Just (7, 205), getFirstTrt = Just {(19, 20), Treatment StandardVaccine}, getFirstFlu = Just {(180, 181), Diagnosis Flu}, getTrtType = Just StandardVaccine, getTimeToFlu = Just 160, getTimeToEndEnr = Just 185}

print $ results !! 3
---> ProcessedSubj {getEnrollment = Just (3, 97), getFirstTrt = Nothing, getFirstFlu = Nothing, getTrtType = Nothing, getTimeToFlu = Nothing, getTimeToEndEnr = Nothing}