{-|
Module      : ExampleFeatures4
Description : Demostrates how to define an outcome monitoring treatment regimes
              over time.
Copyright   : (c) NoviSci, Inc 2020
License     : BSD3
Maintainer  : bsaul@novisci.com
-}

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}

module ExampleFeatures4(
    exampleFeatures4Spec
) where

import Hasklepias
import ExampleEvents ( exampleEvents4 )
import Test.Hspec ( shouldBe, it, Spec, xcontext, describe, pending )

{-
  Example Data and utilities to create such
-}

type EventData a b = (b, a, Text)

t1 :: (a, b, c) -> a
t1 (x , _ , _) = x
t2 :: (a, b, c) -> b
t2 (_ , x , _) = x
t3 :: (a, b, c) -> c
t3 (_ , _ , x) = x

toEvent :: (IntervalSizeable a b, Show a, Integral b) => 
  EventData a b-> Event a
toEvent x = event (beginerval (t1 x) (t2 x))
    (context (UnimplementedDomain ()) (packConcepts [t3 x]))

toEvents :: (Show a, IntervalSizeable a b, Integral b) =>
  [EventData a b] -> Events a
toEvents = sort.map toEvent

sapExample1 :: Events Day
sapExample1 = toEvents [
    (1, fromGregorian 2017 1 1, "index")
  , (1, fromGregorian 2017 3 1, "pcsk")
 ]

sapExample2 :: Events Day
sapExample2 = toEvents [
    (1, fromGregorian 2017 1 1,   "index")
  , (1, fromGregorian 2017 3 1, "wellness")
 ]

p1Events :: Events Int
p1Events = toEvents [
    (1, 1,   "index")
  , (1, 7 + 15, "pcsk")
 ]

p2Events :: Events Int
p2Events = toEvents [
    (1, 1,   "index")
  , (1, 7 + 60, "pcsk")
 ]

p3Events :: Events Int
p3Events = toEvents [
    (1, 1,   "index")
  , (1, 7 + 120, "pcsk")
 ]

p4Events :: Events Int
p4Events = toEvents [
    (1, 1,   "index")
  , (1, 7 + 240, "pcsk")
 ]

p5Events :: Events Int
p5Events = toEvents [
    (1, 1,   "index")
 ]

{-
  Types used for features
-}

data CensorReason =
  -- The order matters here in that if two censoring events occur on the same
  -- day then the reason for censoring will be chosen based on the following 
  -- ordering.
    Death
  | Disenrollment
  | Discontinuation
  | Noncompliance
  | EndOfData
  deriving (Eq, Show, Ord)

instance OccurrenceReason CensorReason

data OutcomeReason =
    Wellness
  | Accident
  -- etc
  deriving (Eq, Show, Ord)

instance OccurrenceReason OutcomeReason

data NegOutcomes b = MkNegOutcome {
    g1 :: CensoredOccurrence CensorReason OutcomeReason b
  , g2 :: CensoredOccurrence CensorReason OutcomeReason b
  , g3 :: CensoredOccurrence CensorReason OutcomeReason b
  , g4 :: CensoredOccurrence CensorReason OutcomeReason b
  , g5 :: CensoredOccurrence CensorReason OutcomeReason b
 } deriving (Eq)

instance (Show b) => Show ( NegOutcomes b ) where
  show (MkNegOutcome x1 x2 x3 x4 x5) =
    "\n g1: " ++ show x1 ++
    "\n g2: " ++ show x2 ++
    "\n g3: " ++ show x3 ++
    "\n g4: " ++ show x4 ++
    "\n g5: " ++ show x5 ++
    "\n"

data ProtocolStatus a =
    Compliant
  | NonCompliant (EventTime a)
  deriving (Eq, Show)

data Protocols a = MkProtocols {
    noInit  :: ProtocolStatus a
  , init30  :: ProtocolStatus a
  , init90  :: ProtocolStatus a
  , init180 :: ProtocolStatus a
  , init365 :: ProtocolStatus a
 } deriving (Eq)

instance (Show a) => Show (Protocols a) where
  show (MkProtocols x1 x2 x3 x4 x5) =
    "\n " ++ show x1 ++
    "\n " ++ show x2 ++
    "\n " ++ show x3 ++
    "\n " ++ show x4 ++
    "\n " ++ show x5 ++
    "\n"

{-
  Helper functions
-}

-- | Duration of follow up in days
followupDuration :: Integral b => b
followupDuration = 365

-- | Duration to not observe events after the begin of index.
washoutDuration :: Integral b => b
washoutDuration = 7

-- | Creates an interval *starting 7 days after the index* and 
--   ending 'followupDuration' days later.
makeFollowupInterval :: ( Integral b
                        , Intervallic i a
                        , IntervalSizeable a b) => 
    b -> Index i a -> Interval a
makeFollowupInterval dur index =  
  beginerval dur (add washoutDuration (begin $ getIndex index))

-- | Creates an interval *starting 7 days after the index* and 
--   ending 'followupDuration' days later.
followupInterval :: (Integral b, IntervalSizeable a b) => 
    Index Interval a -> Interval a
followupInterval = makeFollowupInterval 365

{-
  Functions for defining the study's exposure protocol(s)
-}

protocol :: ( Intervallic i0 a
            , Intervallic i1 a
            , Intervallic i2 a
            , IntervalSizeable a b
            , Filterable container) =>
       (Index i2 a -> i0 a)
    -- ^ Function that maps an index interval to interval during which protocol is evaluated
    -> (i0 a -> container (i1 a) -> ProtocolStatus b)
    -- ^ Function that maps data to a @ProtocolStatus@.
    -> Index i2 a
    -> container (i1 a)
    -> ProtocolStatus b
protocol g f i dat = f (g i) ( filterConcur (g i) dat )

compliantIfNone :: (  IntervalSizeable a b
                    , Intervallic i0 a
                    , Intervallic i1 a
                    , Witherable container ) =>
        i0 a
    -> container (i1 a)
    -> ProtocolStatus b
compliantIfNone i x
  | null x     = Compliant
  | otherwise  = NonCompliant (mkEventTime (fmap (`diff` begin i) (end <$> headMay (toList x))))

compliantIfSome :: ( IntervalSizeable a b
                    , Intervallic i0 a
                    , Intervallic i1 a
                    , Witherable container) =>
       i0 a
    -> container (i1 a)
    -> ProtocolStatus b
compliantIfSome i x
  | null x     = NonCompliant (mkEventTime (Just $ diff (end i) (begin i)))
  | otherwise  = Compliant

protocolNoInit :: ( Integral b
                  , IntervalSizeable a b
                  , Intervallic i0 a
                  , Intervallic i1 a
                  , Witherable container) =>
       Index i0 a
    -> container (i1 a)
    -> ProtocolStatus b
protocolNoInit = protocol (makeFollowupInterval 365) compliantIfNone

protocols :: (  Integral b
              , IntervalSizeable a b
              , Intervallic i0 a
              , Intervallic i1 a
              , Witherable container) =>
       Index i0 a
    -> container (i1 a)
    -> Protocols b
protocols i e = MkProtocols
  ( protocol (makeFollowupInterval 365) compliantIfNone i e)
  ( protocol (makeFollowupInterval 30 ) compliantIfSome i e)
  ( protocol (makeFollowupInterval 90 ) compliantIfSome i e)
  ( protocol (makeFollowupInterval 180) compliantIfSome i e)
  ( protocol (makeFollowupInterval 365) compliantIfSome i e)

-- adminCensor :: (Integral b) => EventTime b -> CensoredOccurrence c o b
-- adminCensor t = MkCensoredOccurrence AdminCensor ( RightCensored t )

compliantOutcome :: (Integral b) => 
     EventTime b
  -> Occurrence OutcomeReason b
  -> Occurrence CensorReason b
  -> CensoredOccurrence CensorReason OutcomeReason b
compliantOutcome
      adminTime
      (MkOccurrence  (oreason, otime))
      (MkOccurrence  (creason, ctime))
  | all (adminTime <) [otime, ctime]   = adminCensor adminTime
  | all (otime <=)    [ctime]          = MkCensoredOccurrence (O oreason) (Uncensored otime)
  | otherwise                          = MkCensoredOccurrence (C creason) (RightCensored ctime)

nonCompliantOutcome :: (Integral b) => 
     EventTime b
  -> EventTime b
  -> Occurrence OutcomeReason b
  -> Occurrence CensorReason b
  -> CensoredOccurrence CensorReason OutcomeReason b
nonCompliantOutcome
      etime
      adminTime
      (MkOccurrence  (oreason, otime))
      (MkOccurrence  (creason, ctime))
  | all (adminTime <) [otime, ctime, etime] = adminCensor adminTime
  | all (otime <=)    [ctime, etime]        = MkCensoredOccurrence (O oreason) (Uncensored otime)
  | etime <= ctime                          = MkCensoredOccurrence (C Noncompliance) (RightCensored etime)
  | otherwise                               = MkCensoredOccurrence (C creason) (RightCensored ctime)

decideOutcome :: (Integral b) =>
     EventTime b      -- ^ admin censoring time
  -> ProtocolStatus b -- ^ pcsk
  -> Occurrence OutcomeReason b    -- ^ time of outcome
  -> Occurrence CensorReason b    -- ^ time of censoring (other than noncompliance)
  -> CensoredOccurrence CensorReason OutcomeReason b
decideOutcome  adminTime exposure outcomeTime censorTime =
  case exposure of
    Compliant      -> compliantOutcome adminTime outcomeTime censorTime
    NonCompliant t -> nonCompliantOutcome t adminTime outcomeTime censorTime

{-
   Features needed to evaluate censoring and outcome events
-}
index :: (Ord a)=> 
  Def (F "events" (Events a) -> F "index" (Index Interval a))
index = defineA (
       makeConceptsFilter ["index"]
    .> intervals
    .> headMay
    .> \case
          Nothing -> makeFeature $ featureDataL ( Other "no index" )
          Just x  -> pure $ makeIndex x
      )

flupEvents :: (Integral b, IntervalSizeable a b) => 
    Def (
     F "index" (Index Interval a)
  -> F "events" (Events a)
  -> F "allFollowupEvents" (Events b))
flupEvents = define (\index es ->
      es
      |> filterConcur ( followupInterval index)
      |> fmap ( diffFromBegin ( followupInterval index ) )
    )

{-
   Censoring Events
-}

death :: Integral b => Def (
     F "allFollowupEvents" (Events b)
  -> F "death" (EventTime b))
death = define (mkEventTime . fmap begin . firstConceptOccurrence ["death"])

disenrollment :: (Integral b, IntervalSizeable a b) =>
  Def (
     F "index" (Index Interval a)
     -- using all events rather than just follow-up events because enrollment
     -- intervals need to be combined first
  -> F "events" (Events a)
  -> F "disenrollment" (EventTime b))
disenrollment = define (\i events ->
  events
  |> makeConceptsFilter ["enrollment"]
  -- combine any concurring enrollment intervals
  |> combineIntervals
  -- find gaps between any enrollment intervals (as well as bounds of followup )
  |> gapsWithin (followupInterval i)
  -- get the first gap longer than 30 days (if it exists)
  |> \x -> (headMay . filter (\x -> duration x > 30)) =<< x
  -- Shift endpoints of intervals so that end of follow up is reference point
  |> fmap (diffFromBegin (followupInterval i))
  -- take the end of this gap as the time of disenrollment
  |> fmap end
  |> mkEventTime )

-- | A collector feature for all censors (except noncompliance)
censorTime :: (Integral b) => Def (
     F "death" (EventTime b)
  -> F "disenrollment" (EventTime b)
  -- etc
  -> F "censortime" (Occurrence CensorReason b))
censorTime = define (
  \dth disrl  ->
      minimum [ makeOccurrence Death dth
              , makeOccurrence Disenrollment disrl
              -- etc
              ]
   )

{- 
  Exposure Definitions
-}

pcskEvents :: Def (
     F "events" (Events a)
  -> F "pcskEvents" (Events a))
pcskEvents = define ( makeConceptsFilter ["pcsk"] )

pcskProtocols :: (Integral b, IntervalSizeable a b) => 
  Def (
     F "index" (Index Interval a)
  -> F "pcskEvents" (Events a)
  -> F "pcskProtocols" (Protocols b) )
pcskProtocols = define protocols

{-
  Outcome definitions
-}

makeg :: (Integral b, IntervalSizeable a b) => 
      b
    -> Index Interval a
    -> ProtocolStatus b
    -> Occurrence OutcomeReason b
    -> Occurrence CensorReason b
    -> CensoredOccurrence CensorReason OutcomeReason b
makeg dur i = decideOutcome 
    (mkEventTime $ Just $ duration (makeFollowupInterval dur i))

makeNegOutcomes :: (Integral b, IntervalSizeable a b) => 
      Index Interval a
   -> Protocols b
   -> Occurrence CensorReason b 
   -> Occurrence OutcomeReason b
   -> NegOutcomes b
makeNegOutcomes i (MkProtocols p1 p2 p3 p4 p5) c o = MkNegOutcome
    (makeg 365 i p1 o c)
    (makeg 30  i p2 o c)
    (makeg 90  i p3 o c)
    (makeg 180 i p4 o c)
    (makeg 365 i p5 o c)

type OutcomeFeature name a b = 
     F "index" (Index Interval a)
  -> F "allFollowupEvents" (Events b)
  -> F "pcskProtocols" (Protocols b)
  -> F "censortime" (Occurrence CensorReason b)
  -> F name (NegOutcomes b)

makeOutcomeDefinition :: 
  ( KnownSymbol name
  , Integral b
  , IntervalSizeable  a b) =>
    [Text]
  -> OutcomeReason
  -> Def (   F "index" (Index Interval a)
          -> F "allFollowupEvents" (Events b)
          -> F "pcskProtocols" (Protocols b)
          -> F "censortime" ( Occurrence CensorReason b) 
          -> F name ( NegOutcomes b))
makeOutcomeDefinition cpt oreason  = define (
    \index events protocols censor  ->
      events
      |> firstConceptOccurrence cpt
      |> \x -> makeOccurrence oreason (mkEventTime (fmap begin x))
      |> makeNegOutcomes index protocols censor
  )

o1 :: (Integral b, IntervalSizeable  a b ) => 
  Def ( OutcomeFeature "wellness" a b)
o1 = makeOutcomeDefinition ["wellness"] Wellness

o2 :: (Integral b, IntervalSizeable  a b ) => 
  Def ( OutcomeFeature "accident" a b)
o2 = makeOutcomeDefinition ["accident"] Accident 

{- 
   Tests of protocols
-}

testProtocols :: (Integral b, IntervalSizeable a b ) => 
  [Event a] -> Feature "pcskProtocols" (Protocols b)
testProtocols input = eval pcskProtocols (idx,  pcev)
  where evs   = pure input
        idx   = eval index evs
        pcev  = eval pcskEvents evs

p1Protocols :: Feature "pcskProtocols" (Protocols Int)
p1Protocols = pure $ MkProtocols
    ( NonCompliant (mkEventTime (Just 15)) )
      Compliant
      Compliant
      Compliant
      Compliant

p2Protocols :: Feature "pcskProtocols" (Protocols Int)
p2Protocols = pure $ MkProtocols
    ( NonCompliant (mkEventTime (Just 60)) )
    ( NonCompliant (mkEventTime (Just 30)) )
      Compliant
      Compliant
      Compliant

p3Protocols :: Feature "pcskProtocols" (Protocols Int)
p3Protocols = pure $ MkProtocols
    ( NonCompliant (mkEventTime (Just 120)) )
    ( NonCompliant (mkEventTime (Just 30)) )
    ( NonCompliant (mkEventTime (Just 90)) )
      Compliant
      Compliant

p4Protocols :: Feature "pcskProtocols" (Protocols Int)
p4Protocols = pure $ MkProtocols
    ( NonCompliant (mkEventTime (Just 240)) )
    ( NonCompliant (mkEventTime (Just 30)) )
    ( NonCompliant (mkEventTime (Just 90)) )
    ( NonCompliant (mkEventTime (Just 180)) )
      Compliant

p5Protocols :: Feature "pcskProtocols" (Protocols Int)
p5Protocols = pure $ MkProtocols
      Compliant
    ( NonCompliant (mkEventTime (Just 30)) )
    ( NonCompliant (mkEventTime (Just 90)) )
    ( NonCompliant (mkEventTime (Just 180)) )
    ( NonCompliant (mkEventTime (Just 365)) )

{- 
   Tests of outcomes
-}

testOutcomes :: ( Integral b, IntervalSizeable  a b ) => 
      [Event a]
   -> (Feature "wellness" ( NegOutcomes b ), Feature "accident" (NegOutcomes b) )
testOutcomes input = (
      eval o1 (idx, flevs, prot, ctime) 
    , eval o2 (idx, flevs, prot, ctime) )
  where evs   = pure input
        idx   = eval index evs
        flevs = eval flupEvents (idx, evs)
        pcev  = eval pcskEvents evs
        dth   = eval death flevs
        disen = eval disenrollment (idx, evs)
        prot  = eval pcskProtocols (idx, pcev)
        ctime = eval censorTime (dth, disen)

p1Outcomes :: (Integral b) => 
  ( Feature "wellness" ( NegOutcomes b)
  , Feature "accident" ( NegOutcomes b)) 
p1Outcomes = (
    pure $ MkNegOutcome 
      (MkCensoredOccurrence (C Noncompliance) (RightCensored (mkEventTime (Just 15))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 30))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 90))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 180))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 365))))
  , pure $ MkNegOutcome 
      (MkCensoredOccurrence (C Noncompliance) (RightCensored (mkEventTime (Just 15))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 30))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 90))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 180))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 365)))) 
 )


p1Outcomes' :: (Integral b) => 
  ( Feature "wellness" (NegOutcomes b)
  , Feature "accident" (NegOutcomes b)) 
p1Outcomes' = (
    pure $ MkNegOutcome 
      (MkCensoredOccurrence (C Noncompliance) (RightCensored (mkEventTime (Just 15))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 30))))
      (MkCensoredOccurrence (O Wellness)      (Uncensored (mkEventTime (Just 51))))
      (MkCensoredOccurrence (O Wellness)      (Uncensored (mkEventTime (Just 51))))
      (MkCensoredOccurrence (O Wellness)      (Uncensored (mkEventTime (Just 51))))
  , pure $ MkNegOutcome 
      (MkCensoredOccurrence (C Noncompliance) (RightCensored (mkEventTime (Just 15))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 30))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 90))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 180))))
      (MkCensoredOccurrence  AdminCensor      (RightCensored (mkEventTime (Just 365)))) 
 )

{-
  Test specs
-}

exampleFeatures4Spec :: Spec
exampleFeatures4Spec = do
    describe "tests of exposure protocols" $
      do
        it "p1" $
          testProtocols p1Events `shouldBe` p1Protocols
        it "p2" $
          testProtocols p2Events `shouldBe` p2Protocols
        it "p3" $
          testProtocols p3Events `shouldBe` p3Protocols
        it "p4" $
          testProtocols p4Events `shouldBe` p4Protocols
        it "p5" $
          testProtocols p5Events `shouldBe` p5Protocols

    describe "tests of outcomes" $
      do
        it "p1" $
          testOutcomes p1Events `shouldBe` p1Outcomes
        it "p1'" $
          testOutcomes (sort p1Events <> [toEvent (1, 59, "wellness")] ) 
            `shouldBe` p1Outcomes'

    -- describe "SAP examples" $ 
    --   do
    --     it "sap example 1" pending
        -- testOutcomes sapExample1 `shouldBe` ???