module classroom_member

  use :: iso_c_binding, only : c_double, c_int, c_int64_t
  use :: prob
  use :: dists
  use :: classroom_params
  implicit none

  private
  integer, parameter :: dp = c_double
  integer, parameter :: i4 = c_int
  integer, parameter :: i8 = c_int64_t

  type, public :: ClassMember
    integer(kind=c_int) :: cohort=0 ! 1:pupil, 2:teacher, 3:TA, 4:nonclass
    ! Disease progression
    logical :: susceptible = .False.
    logical :: exposed = .False.
    logical :: unwell = .False.
    logical :: recovered = .False.

    ! Social status
    logical :: quarantined = .False.
    logical :: in_class = .True.
    logical :: immune = .False.

    ! Test status
    logical :: rapid_test_positive = .False.
    logical :: pcr_test_positive = .False.
    integer(kind=c_int) :: rapid_positive_day = 0

    ! Individual's properties
    integer(kind=c_int) :: days_infected = 0
    integer(kind=c_int) :: days_unwell = 0
    integer(kind=c_int) :: unwell_time
    integer(kind=c_int) :: days_quarantined = 0
    integer(kind=c_int) :: incubation_time = 0
    integer(kind=c_int) :: tost_time = 0
    integer(kind=c_int) :: infection_clearing_time ! = incubation_time + unwell_time.
    logical :: symptomatic = .False.
    logical :: infected = .False.
    logical :: vaccinated = .False.
    integer(kind=c_int) :: daily_contacts = 3
    real(kind=c_double) :: mitigation_factor = 1.0_dp
    integer(kind=c_int) :: reduced_daily_contacts = 3
    type(relative_transmissibilty) :: rel_trans
    real(kind=c_double) :: transmission_prob = 0.035_dp
  ! contains
    ! procedure, public, pass(self) :: describe => describe_ClassMember
    ! procedure, public, pass(self) :: advance => advanceSEIR
  end type ClassMember

  type, public, extends(ClassMember) :: Pupil
    ! Nothing new
  end type Pupil

  interface Pupil
    module procedure init_Pupil_params
  end interface Pupil

  type, public, extends(ClassMember) :: Teacher
    logical :: temporary = .False.
  end type Teacher

  interface Teacher
    module procedure init_Teacher_params
  end interface Teacher

  ! interface ClassMember
  !   module procedure init_ClassMember_params
  ! end interface ClassMember

contains

function init_Pupil_params(Params, infected, symptomatic, vaccinated) result(self)
    implicit none
    type(Pupil) :: self
    type(Parameters), intent(in) :: Params
    logical, intent(in) :: infected
    logical, intent(in) :: symptomatic
    logical, intent(in) :: vaccinated
    integer(kind=c_int) :: j

    real(kind=dp) :: sample_contacts

    ! Set cohort of member
    self%cohort = 1

    ! Set whether member is vaccinated -- likely comes from a Bernoulli draw
    self%vaccinated = vaccinated
    ! Set whether member is infected -- likely comes from a Bernoulli draw
    self%infected = infected
    ! Set whether member will be symptomatic -- likely comes from a Bernoulli draw
    self%symptomatic = symptomatic
    
    ! Set transmission probability
    if (symptomatic) then
      self%transmission_prob = Params%Transmission_prob_sympt_child
    else
      self%transmission_prob = Params%Transmission_prob_asympt_child
    end if

    if (vaccinated) then
      self%transmission_prob = self%transmission_prob * Params%Vaccine_Transmission_prob_factor
    end if
    
    ! Initiate number of days isolated
    self%days_quarantined = 0

    ! Set incubation time -- drawn from lognormal distribution pdf (McAloon et al 2020)
    self%incubation_time =  nint(Params%incubation_time%sample(), kind=c_int)

    ! Set maximum days unwell -- note there is an unwell time asymptomatics too -- TODO : Should be draw from pdf
    self%unwell_time = nint(Params%unwell_time_child%sample(), kind=c_int)

    ! Set infection clearing time = incubation time+unwell time
    self%infection_clearing_time = self%incubation_time + self%unwell_time

    ! Set tost (time from onset of symptoms to transmission) -- used to set variation in transmissibilty over time (Ferretti et al 2020)
    ! Note: even asymptomatics are given a tost, to use in the transmissibility function
    self%tost_time = nint(Params%tost%sample(), kind=c_int)

    self%rel_trans = relative_transmissibilty(self%incubation_time)

    ! Set number of daily contacts -- comes from a draw from an empirical pdf
    sample_contacts = Params%pupil_contacts%sample()
    self%daily_contacts = nint(sample_contacts, kind=c_int)

    ! If infected, how long have they been infected?
    ! infection started on day -days_infected
    ! Sample from Uniform(0, infection_clearing_time)
    if (infected) then
      self%days_infected = int(pdists%randint(0,self%infection_clearing_time), kind=c_int)
    else
      self%days_infected = 0
    end if
    
    self%mitigation_factor = Params%mitigations%pupil_mitigation_factor
    
    ! Reduce the number of daily contacts based on mitigation factor
    self%reduced_daily_contacts = nint(self%mitigation_factor * self%daily_contacts, kind=c_int)

    ! Initialize place in disease progression (SEUQR) model
    if (.not. self%infected) then
      ! Member is not infected, so is in Susceptible compartment
      self%susceptible = .True.
      ! and is also in the class
      self%in_class = .True.
    else
      ! Member is infected, so is not in Susceptible compartment
      self%susceptible = .False.
      ! but the compartment depends on how long they have been infected and whether they are symptomatic
      if (self%symptomatic) then
        
        if (self%days_infected <= self%incubation_time) then
          ! infected member is in the incubation phase, so in the Exposed compartment
          self%exposed = .True.
          ! and not currently unwell
          self%unwell = .False.
          ! but is in the class
          self%in_class = .True.
        else
          ! infected member has passed through the incubation phase.
          self%exposed = .False.

          ! Are they currently unwell?
          if (self%days_infected <= self%infection_clearing_time) then
            ! The member is currently unwell
            self%unwell = .True.
            self%days_unwell = self%days_infected-self%incubation_time
          else
            ! The member is recovered
            self%recovered = .True.
            self%days_unwell = self%unwell_time
            self%immune = .True.
          end if

          ! Is the member in quarantine?
          ! First work out on which day the symptoms were detected
          do j=-self%days_infected+self%incubation_time, 0
            self%quarantined = pdists%bernoulli_rvs(Params%Detection_prob_sympt)
            if (self%quarantined) then
              self%days_quarantined = -j
              exit
            end if
          end do
          if (self%days_quarantined > Params%quarantine_days) then
            self%quarantined = .False.
            self%in_class = .True.
          else
            self%quarantined = .True.
            self%in_class = .False.
          end if
        end if
      else
        ! If the member is asymptomatic, then they would not be isolated
        self%exposed = .True.
        self%unwell = .False.
      end if
    end if

    return
  end function init_Pupil_params

  function init_Teacher_params(Params, infected, symptomatic, vaccinated, temporary) result(self)
    implicit none
    type(Teacher) :: self
    type(Parameters), intent(in) :: Params
    logical, intent(in) :: infected
    logical, intent(in) :: symptomatic
    logical, intent(in) :: vaccinated
    logical, intent(in), optional :: temporary
    integer(kind=c_int) :: j

    real(kind=dp) :: sample_contacts

    ! Set cohort of member
    if (present(temporary)) then
      if (temporary) then
        self%cohort = 5
        self%temporary = .True.
      else
        self%cohort = 2
        self%temporary = .False.
      end if
    else
      self%cohort = 2
      self%temporary = .False.
    end if

    ! Set whether member is vaccinated -- likely comes from a Bernoulli draw
    self%vaccinated = vaccinated

    ! Set whether member is infected -- likely comes from a Bernoulli draw
    self%infected = infected

    ! Set whether member will be symptomatic -- likely comes from a Bernoulli draw
    self%symptomatic = symptomatic

    ! Set transmission probability
    if (symptomatic) then
      self%transmission_prob = Params%Transmission_prob_sympt_adult
    else
      self%transmission_prob = Params%Transmission_prob_asympt_adult
    end if

    if (vaccinated) then
      self%transmission_prob = self%transmission_prob * Params%Vaccine_Transmission_prob_factor
    end if
    
    ! Initiate number of days isolated
    self%days_quarantined = 0

    ! Set incubation time -- drawn from lognormal distribution pdf (McAloon et al 2020)
    self%incubation_time =  nint(Params%incubation_time%sample(), kind=c_int)

    ! Set maximum days unwell -- note there is an unwell time asymptomatics too
    self%unwell_time = nint(Params%unwell_time_adult%sample(), kind=c_int)

    ! Set infection clearing time = incubation time+unwell time
    self%infection_clearing_time = self%incubation_time + self%unwell_time

    ! Set tost (time from onset of symptoms to transmission) -- used to set variation in transmissibilty over time (Ferretti et al 2020)
    ! Note: even asymptomatics are given a tost, to use in the transmissibility function
    self%tost_time = nint(Params%tost%sample(), kind=c_int)

    self%rel_trans = relative_transmissibilty(self%incubation_time)

    ! Set number of daily contacts -- comes from a draw from an empirical pdf
    sample_contacts = Params%teacher_contacts%sample()
    self%daily_contacts = nint(sample_contacts, kind=c_int)

    ! If infected, how long have they been infected?
    ! infection started on day -days_infected
    
    ! Get the days infected for the teacher.  Depends whether they are symptomatic.
      ! TODO : enhance teacher -- replaced by temporary teacher
    if (infected) then
      if (symptomatic) then
        ! If the teacher is symptomatic, the days_infected can only be up to the time when symptoms are detected
        ! for simplicity here, this is set to the incubation_time, but should be fixed
        self%days_infected = int(pdists%randint(0,self%incubation_time), kind=c_int)  !!! TODO : IMPORTANT: set this correctly
      else
        ! If the teacher is asymptomatic, the days_infected can be up to infection_clearing_time
        self%days_infected = int(pdists%randint(0,self%infection_clearing_time), kind=c_int) !!! IMPORTANT: set this correctly
      end if
    else
      self%days_infected = 0
    end if

    self%mitigation_factor = Params%mitigations%teacher_mitigation_factor

    ! Reduce the number of daily contacts based on mitigation factor
    self%reduced_daily_contacts = nint(self%mitigation_factor * self%daily_contacts, kind=c_int)

    ! Initialize place in disease progression (SEUQR) model
    if (.not. self%infected) then
      ! Member is not infected, so is in Susceptible compartment
      self%susceptible = .True.
      ! and is also in the class
      self%in_class = .True.
    else
      ! Member is infected, so is not in Susceptible compartment
      self%susceptible = .False.
      ! but the compartment depends on how long they have been infected and whether they are symptomatic
      if (self%symptomatic) then
        
        if (self%days_infected <= self%incubation_time) then
          ! infected member is in the incubation phase, so in the Exposed compartment
          self%exposed = .True.
          ! and not currently unwell
          self%unwell = .False.
          ! but is in the class
          self%in_class = .True.
        else
          ! infected member has passed through the incubation phase.
          self%exposed = .False.

          ! Are they currently unwell?
          if (self%days_infected <= self%infection_clearing_time) then
            ! The member is currently unwell
            self%unwell = .True.
            self%days_unwell = self%days_infected-self%incubation_time
          else
            ! The member is recovered
            self%recovered = .True.
            self%days_unwell = self%unwell_time
            self%immune = .True.
          end if

          ! Is the member in quarantine?
          ! First work out on which day the symptoms were detected
          do j=-self%days_infected+self%incubation_time, 0
            self%quarantined = pdists%bernoulli_rvs(Params%Detection_prob_sympt)
            if (self%quarantined) then
              self%days_quarantined = -j
              exit
            end if
          end do
          if (self%days_quarantined > Params%quarantine_days) then
            self%quarantined = .False.
            self%in_class = .True.
          else
            self%quarantined = .True.
            self%in_class = .False.
          end if
        end if
      else
        ! If the member is asymptomatic, then they would not be isolated
        self%exposed = .True.
        self%unwell = .False.
      end if
    end if

    return

  end function init_Teacher_params

  ! function init_ClassMember_params(Params, cohort, infected, symptomatic) result(self)
  !   implicit none
  !   type(ClassMember) :: self
  !   type(Parameters), intent(in) :: Params
  !   integer(kind=c_int), intent(in) :: cohort
  !   logical, intent(in) :: infected
  !   logical, intent(in) :: symptomatic
  !   integer(kind=c_int) :: j

  !   ! Set cohort of member
  !   self%cohort = cohort
  !   if (cohort==5) self%temporary = .True.

  !   ! Set whether member is infected -- likely comes from a Bernoulli draw
  !   self%infected = infected
  !   ! Set whether member will be symptomatic -- likely comes from a Bernoulli draw
  !   self%symptomatic = symptomatic
  !   if (symptomatic) then
  !     self%transmission_prob = Params%Transmission_prob_sympt
  !   else
  !     self%transmission_prob = Params%Transmission_prob_asympt
  !   end if
    
  !   ! Initiate number of days isolated
  !   self%days_quarantined = 0

  !   ! Set incubation time -- drawn from lognormal distribution pdf (McAloon et al 2020)
  !   self%incubation_time =  nint(Params%incubation_time%sample(), kind=c_int)

  !   ! Set maximum days unwell -- note there is an unwell time asymptomatics too -- TODO : Should be draw from pdf
  !   self%unwell_time = Params%unwell_time !! TODO : change to draw from suitable pdf

  !   ! Set infection clearing time = incubation time+unwell time
  !   self%infection_clearing_time = self%incubation_time + self%unwell_time

  !   ! Set tost (time from onset of symptoms to transmission) -- used to set variation in transmissibilty over time (Ferretti et al 2020)
  !   ! Note: even asymptomatics are given a tost, to use in the transmissibility function
  !   self%tost_time = nint(Params%tost%sample(), kind=c_int)

  !   self%rel_trans = relative_transmissibilty(self%incubation_time)

  !   ! Set number of daily contacts -- likely comes from a draw from a pdf
  !   if (cohort==1) then
  !     self%daily_contacts = nint(pdists%loglogistic_rvs(shape=pupil_contact_params%shape, &
  !                                                     scale=pupil_contact_params%scale), &
  !                                kind=c_int)
  !   else
  !     self%daily_contacts = nint(pdists%loglogistic_rvs(shape=teacher_contact_params%shape, &
  !                                                     scale=teacher_contact_params%scale), &
  !                                kind=c_int)
  !   end if

  !   ! If infected, how long have they been infected?
  !   ! infection started on day -days_infected
  !   ! If it's not a teacher, we can sample from Uniform(0, infection_clearing_time)
  !   if (cohort .ne. 2) then
  !     if (infected) then
  !       self%days_infected = int(pdists%randint(0,self%infection_clearing_time), kind=c_int)
  !     else
  !       self%days_infected = 0
  !     end if
  !   else
  !     ! Get the days infected for the teacher.  Depends whether they are symptomatic.
  !       ! TODO : enhance teacher -- replaced by temporary teacher
  !     if (infected) then
  !       if (symptomatic) then
  !         ! If the teacher is symptomatic, the days_infected can only be up to the time when symptoms are detected
  !         ! for simplicity here, this is set to the incubation_time, but should be fixed
  !         self%days_infected = int(pdists%randint(0,self%incubation_time), kind=c_int)  !!! TODO : IMPORTANT: set this correctly
  !       else
  !         ! If the teacher is asymptomatic, the days_infected can be up to infection_clearing_time
  !         self%days_infected = int(pdists%randint(0,self%infection_clearing_time), kind=c_int) !!! IMPORTANT: set this correctly
  !       end if
  !     else
  !       self%days_infected = 0
  !     end if
  !   end if

  !   ! Determine the mitigation factor based on cohort
  !   if (cohort==1) then
  !     self%mitigation_factor = Params%pupil_mitigation_factor
  !   else if (cohort==2) then
  !     self%mitigation_factor = Params%teacher_mitigation_factor
  !   else if (cohort==3) then
  !     self%mitigation_factor = Params%teacher_mitigation_factor
  !   else if (cohort==4) then
  !     self%mitigation_factor = Params%nonclass_mitigation_factor
  !   end if

  !   ! Reduce the number of daily contacts based on mitigation factor
  !   self%reduced_daily_contacts = nint(self%mitigation_factor * self%daily_contacts, kind=c_int)

  !   ! Initialize place in disease progression (SEUQR) model
  !   if (.not. self%infected) then
  !     ! Member is not infected, so is in Susceptible compartment
  !     self%susceptible = .True.
  !     ! and is also in the class
  !     self%in_class = .True.
  !   else
  !     ! Member is infected, so is not in Susceptible compartment
  !     self%susceptible = .False.
  !     ! but the compartment depends on how long they have been infected and whether they are symptomatic
  !     if (self%symptomatic) then
        
  !       if (self%days_infected <= self%incubation_time) then
  !         ! infected member is in the incubation phase, so in the Exposed compartment
  !         self%exposed = .True.
  !         ! and not currently unwell
  !         self%unwell = .False.
  !         ! but is in the class
  !         self%in_class = .True.
  !       else
  !         ! infected member has passed through the incubation phase.
  !         self%exposed = .False.

  !         ! Are they currently unwell?
  !         if (self%days_infected <= self%infection_clearing_time) then
  !           ! The member is currently unwell
  !           self%unwell = .True.
  !           self%days_unwell = self%days_infected-self%incubation_time
  !         else
  !           ! The member is recovered
  !           self%recovered = .True.
  !           self%days_unwell = self%unwell_time
  !           self%immune = .True.
  !         end if

  !         ! Is the member in quarantine?
  !         ! First work out on which day the symptoms were detected
  !         do j=-self%days_infected+self%incubation_time, 0
  !           self%quarantined = pdists%bernoulli_rvs(Params%Detection_prob_sympt)
  !           if (self%quarantined) then
  !             self%days_quarantined = -j
  !             exit
  !           end if
  !         end do
  !         if (self%days_quarantined > Params%quarantine_days) then
  !           self%quarantined = .False.
  !           self%in_class = .True.
  !         else
  !           self%quarantined = .True.
  !           self%in_class = .False.
  !         end if
  !       end if
  !     else
  !       ! If the member is asymptomatic, then they would not be isolated
  !       self%exposed = .True.
  !       self%unwell = .False.
  !     end if
  !   end if

  !   return

  ! end function init_ClassMember_params

  subroutine describe_ClassMember(self)
    implicit none
    class(ClassMember), intent(in) :: self

    character(len=:), allocatable :: SEUR

    if (self%susceptible) then
      SEUR = 'Susceptible'
    elseif (self%exposed) then
      SEUR = 'Exposed'
    elseif (self%unwell) then
      SEUR = 'Unwell'
    elseif (self%recovered) then
      SEUR = 'Recovered'
    else
      SEUR = '**** UNSPECIFIED ****'
    end if

    write(*,*) 'Cohort = ', self%cohort
    write(*,*) 'Disease stage = ', SEUR
    write(*,*) ' Quarantined = ', self%quarantined
    write(*,*) 'infected = ', self%infected
    write(*,*) 'days_infected = ', self%days_infected
    write(*,*) 'incubation_time = ', self%incubation_time
    write(*,*) 'symptomatic = ', self%symptomatic
    write(*,*) 'daily_contacts = ', self%daily_contacts
    write(*,*) 'mitigation_factor = ', self%mitigation_factor
    write(*,*) 'reduced_daily_contacts = ', self%reduced_daily_contacts
    write(*,*) 'in_class = ', self%in_class
    write(*,*) '----------------------'
    return
  end subroutine describe_ClassMember

end module classroom_member