{-# LANGUAGE DeriveGeneric #-}

-- |
-- Module: Epidemic.Types.Population
-- Copyright: (c) 2021 Alexander E. Zarebski
-- License: MIT
--
-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>
-- Stability: unstable
-- Portability: ghc
--
-- This module defines some types and functions for working with identifiers,
-- persons, people, and populations.
--
--   * An 'Identifier' is used as a unique label for a 'Person',
--   * a 'Person' is a single individual,
--   * a group of 'People' is a collection of persons,
--   * and 'Population' is a typeclass for working with people that have some structure.
--

module Epidemic.Types.Population
  ( Person(Person)
  , People(People)
  , Population(..)
  , Identifier(Identifier)
  , asPeople
  , includesPerson
  , haveCommonPeople
  , nullPeople
  , numPeople
  , addPerson
  , removePerson
  , personByteString
  ) where

import qualified Data.Aeson              as Json
import qualified Data.ByteString.Builder as BBuilder
import qualified Data.Vector             as V
import           GHC.Generics

-- | Class of types that can represent populations in an epidemic simulation.
class Population a where
  susceptiblePeople :: a -> Maybe People
  infectiousPeople :: a -> Maybe People
  removedPeople :: a -> Maybe People
  isInfected :: a -> Bool

-- | A type to hold an integer which is unique to each 'Person'.
newtype Identifier =
  Identifier Integer
  deriving (Int -> Identifier -> ShowS
[Identifier] -> ShowS
Identifier -> String
(Int -> Identifier -> ShowS)
-> (Identifier -> String)
-> ([Identifier] -> ShowS)
-> Show Identifier
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Identifier] -> ShowS
$cshowList :: [Identifier] -> ShowS
show :: Identifier -> String
$cshow :: Identifier -> String
showsPrec :: Int -> Identifier -> ShowS
$cshowsPrec :: Int -> Identifier -> ShowS
Show, (forall x. Identifier -> Rep Identifier x)
-> (forall x. Rep Identifier x -> Identifier) -> Generic Identifier
forall x. Rep Identifier x -> Identifier
forall x. Identifier -> Rep Identifier x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Identifier x -> Identifier
$cfrom :: forall x. Identifier -> Rep Identifier x
Generic, Identifier -> Identifier -> Bool
(Identifier -> Identifier -> Bool)
-> (Identifier -> Identifier -> Bool) -> Eq Identifier
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Identifier -> Identifier -> Bool
$c/= :: Identifier -> Identifier -> Bool
== :: Identifier -> Identifier -> Bool
$c== :: Identifier -> Identifier -> Bool
Eq)

instance Json.FromJSON Identifier

instance Json.ToJSON Identifier

-- | A type to represent a single person in a group of 'People'
newtype Person =
  Person Identifier
  deriving (Int -> Person -> ShowS
[Person] -> ShowS
Person -> String
(Int -> Person -> ShowS)
-> (Person -> String) -> ([Person] -> ShowS) -> Show Person
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Person] -> ShowS
$cshowList :: [Person] -> ShowS
show :: Person -> String
$cshow :: Person -> String
showsPrec :: Int -> Person -> ShowS
$cshowsPrec :: Int -> Person -> ShowS
Show, (forall x. Person -> Rep Person x)
-> (forall x. Rep Person x -> Person) -> Generic Person
forall x. Rep Person x -> Person
forall x. Person -> Rep Person x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Person x -> Person
$cfrom :: forall x. Person -> Rep Person x
Generic, Person -> Person -> Bool
(Person -> Person -> Bool)
-> (Person -> Person -> Bool) -> Eq Person
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Person -> Person -> Bool
$c/= :: Person -> Person -> Bool
== :: Person -> Person -> Bool
$c== :: Person -> Person -> Bool
Eq)

instance Json.FromJSON Person

instance Json.ToJSON Person

-- | A type to represent a population.
newtype People =
  People (V.Vector Person)
  deriving (Int -> People -> ShowS
[People] -> ShowS
People -> String
(Int -> People -> ShowS)
-> (People -> String) -> ([People] -> ShowS) -> Show People
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [People] -> ShowS
$cshowList :: [People] -> ShowS
show :: People -> String
$cshow :: People -> String
showsPrec :: Int -> People -> ShowS
$cshowsPrec :: Int -> People -> ShowS
Show, People -> People -> Bool
(People -> People -> Bool)
-> (People -> People -> Bool) -> Eq People
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: People -> People -> Bool
$c/= :: People -> People -> Bool
== :: People -> People -> Bool
$c== :: People -> People -> Bool
Eq, (forall x. People -> Rep People x)
-> (forall x. Rep People x -> People) -> Generic People
forall x. Rep People x -> People
forall x. People -> Rep People x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep People x -> People
$cfrom :: forall x. People -> Rep People x
Generic)

instance Json.FromJSON People

instance Json.ToJSON People

-- | A list of persons as a people
asPeople :: [Person] -> People
asPeople :: [Person] -> People
asPeople [Person]
persons = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ [Person] -> Vector Person
forall a. [a] -> Vector a
V.fromList [Person]
persons

-- | Predicate for whether a person is one of the people
includesPerson :: People -> Person -> Bool
includesPerson :: People -> Person -> Bool
includesPerson (People Vector Person
persons) Person
person = Person -> Vector Person -> Bool
forall a. Eq a => a -> Vector a -> Bool
V.elem Person
person Vector Person
persons

-- | Predicate for whether two sets of people have any members in common.
haveCommonPeople :: People -> People -> Bool
haveCommonPeople :: People -> People -> Bool
haveCommonPeople (People Vector Person
ps1) (People Vector Person
ps2) = (Person -> Bool) -> Vector Person -> Bool
forall a. (a -> Bool) -> Vector a -> Bool
V.any (Person -> Vector Person -> Bool
forall a. Eq a => a -> Vector a -> Bool
`V.elem` Vector Person
ps2) Vector Person
ps1

-- | Predicate for whether there are any people
nullPeople :: People -> Bool
nullPeople :: People -> Bool
nullPeople (People Vector Person
persons) = Vector Person -> Bool
forall a. Vector a -> Bool
V.null Vector Person
persons

-- | The number of people
numPeople :: People -> Int
numPeople :: People -> Int
numPeople (People Vector Person
persons) = Vector Person -> Int
forall a. Vector a -> Int
V.length Vector Person
persons

-- | Add a person to a group of people
addPerson :: Person -> People -> People
addPerson :: Person -> People -> People
addPerson Person
person (People Vector Person
persons) = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ Person -> Vector Person -> Vector Person
forall a. a -> Vector a -> Vector a
V.cons Person
person Vector Person
persons

-- | Remove a person from a group of people
removePerson :: Person -> People -> People
removePerson :: Person -> People -> People
removePerson Person
person (People Vector Person
persons) = Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ (Person -> Bool) -> Vector Person -> Vector Person
forall a. (a -> Bool) -> Vector a -> Vector a
V.filter (Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
/= Person
person) Vector Person
persons

-- | A bytestring builder for a person
personByteString :: Person -> BBuilder.Builder
personByteString :: Person -> Builder
personByteString (Person (Identifier Integer
n)) = Integer -> Builder
BBuilder.integerDec Integer
n