module classroom_records

  use :: iso_c_binding, only : c_double, c_int, c_int64_t
  use :: utilities, only : push
  use :: classroom_params

  implicit none

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

  type, public :: member_record
    integer(kind=c_int64_t), dimension(:), allocatable :: susceptible
    integer(kind=c_int64_t), dimension(:), allocatable :: exposed
    integer(kind=c_int64_t), dimension(:), allocatable :: unwell
    integer(kind=c_int64_t), dimension(:), allocatable :: quarantined
    integer(kind=c_int64_t), dimension(:), allocatable :: recovered
    integer(kind=c_int64_t), dimension(:), allocatable :: infected
    integer(kind=c_int64_t), dimension(:), allocatable :: inclass
  end type member_record

  type, public :: contact_record
    integer(kind=c_int), dimension(:), allocatable :: time_step
    integer(kind=c_int64_t), dimension(:), allocatable :: indxA
    integer(kind=c_int64_t), dimension(:), allocatable :: indxB
    integer(kind=c_int), dimension(:), allocatable :: transmission
  end type contact_record

  type, public :: Records
    integer(kind=c_int), dimension(:), allocatable :: num_infected_pupils
    integer(kind=c_int), dimension(:), allocatable :: num_in_class
    integer(kind=c_int), dimension(:), allocatable :: infected_teacher
    integer(kind=c_int), dimension(:), allocatable :: susceptible
    integer(kind=c_int), dimension(:), allocatable :: exposed
    integer(kind=c_int), dimension(:), allocatable :: unwell
    integer(kind=c_int), dimension(:), allocatable :: quarantined
    integer(kind=c_int), dimension(:), allocatable :: recovered

    integer(kind=c_int) :: total_infected_pupils
    integer(kind=c_int) :: total_infected

    type(contact_record) :: contacts
    type(member_record), dimension(:), allocatable :: cat_lists

    integer(kind=c_int64_t), dimension(:), allocatable :: secondary_infections
    integer(kind=c_int64_t), dimension(:), allocatable :: pupil_absent_days

    integer(kind=c_int) :: num_classroom_transmission = 0
    integer(kind=c_int) :: num_community_transmission = 0
  contains
    procedure, pass(self) :: add_contact_record
    procedure, pass(self) :: copy => copyRecords
    final :: destroy_Records
  end type Records

  interface Records
    module procedure init_Records
  end interface Records

contains
  
  function init_Records(days, class_size, num_in_class, &
                        num_infected, num_infected_pupils, infected_teacher, &
                        susceptible, exposed, unwell, quarantined, recovered,&
                        num_community_transmission) result(self)
    implicit none
    type(Records) :: self
    integer(kind=c_int), intent(in) :: days
    integer(kind=c_int), intent(in) :: class_size
    integer(kind=c_int), intent(in) :: num_in_class
    integer(kind=c_int), intent(in) :: num_infected
    integer(kind=c_int), intent(in) :: num_infected_pupils
    logical, intent(in) :: infected_teacher
    integer(kind=c_int), intent(in) :: susceptible
    integer(kind=c_int), intent(in) :: exposed
    integer(kind=c_int), intent(in) :: unwell
    integer(kind=c_int), intent(in) :: quarantined
    integer(kind=c_int), intent(in) :: recovered
    integer(kind=c_int), intent(in) :: num_community_transmission

    if (allocated(self%num_in_class)) deallocate(self%num_in_class)
    if (allocated(self%num_infected_pupils)) deallocate(self%num_infected_pupils)
    if (allocated(self%infected_teacher)) deallocate(self%infected_teacher)
    if (allocated(self%susceptible)) deallocate(self%susceptible)
    if (allocated(self%exposed)) deallocate(self%exposed)
    if (allocated(self%unwell)) deallocate(self%unwell)
    if (allocated(self%quarantined)) deallocate(self%quarantined)
    if (allocated(self%recovered)) deallocate(self%recovered)
    if (allocated(self%infected_teacher)) deallocate(self%infected_teacher)

    if (allocated(self%contacts%time_step)) deallocate(self%contacts%time_step)
    if (allocated(self%contacts%indxA)) deallocate(self%contacts%indxA)
    if (allocated(self%contacts%indxB)) deallocate(self%contacts%indxB)
    if (allocated(self%contacts%transmission)) deallocate(self%contacts%transmission)

    if (allocated(self%cat_lists)) deallocate(self%cat_lists)

    if (allocated(self%secondary_infections)) deallocate(self%secondary_infections)
    if (allocated(self%pupil_absent_days)) deallocate(self%pupil_absent_days)

    allocate(self%num_in_class(1))
    allocate(self%num_infected_pupils(1))
    allocate(self%infected_teacher(1))
    allocate(self%susceptible(1))
    allocate(self%exposed(1))
    allocate(self%unwell(1))
    allocate(self%quarantined(1))
    allocate(self%recovered(1))
    self%num_in_class(1) = num_in_class
    self%num_infected_pupils(1) = num_infected_pupils
    self%infected_teacher(1) = merge(1_i4, 0_i4, infected_teacher) ! 1 if infected_teacher, else 0

    self%susceptible(1) = susceptible
    self%exposed(1) = exposed
    self%unwell(1) = unwell
    self%quarantined(1) = quarantined
    self%recovered(1) = recovered

    self%num_classroom_transmission = 0
    self%num_community_transmission = num_community_transmission

    self%total_infected_pupils = num_infected_pupils
    
    self%total_infected = num_infected

    allocate(self%cat_lists(days))

    allocate(self%secondary_infections(class_size+2))
    self%secondary_infections(:) = -1_i8

    allocate(self%pupil_absent_days(class_size))
    self%pupil_absent_days(:) = 0_i8

    return
  end function init_Records

  subroutine add_contact_record(self, time_step, &
                                indxA, indxB, &
                                transmission)
    implicit none
    class(Records), intent(inout) :: self
    integer(kind=c_int), intent(in) :: time_step
    integer(kind=c_int64_t), intent(in) :: indxA
    integer(kind=c_int64_t), intent(in) :: indxB
    logical, intent(in), optional :: transmission

    integer(kind=c_int) :: transmission_
    
    if (present(transmission)) then
      transmission_ = merge(1_i4, 0_i4, transmission)
    else
      transmission_ = 0_i4
    end if
    
    call push(self%contacts%time_step, time_step)
    call push(self%contacts%indxA, indxA)
    call push(self%contacts%indxB, indxB)
    call push(self%contacts%transmission, transmission_)
    
    return
  end subroutine add_contact_record

  function copyRecords(self) result(new_record)
    implicit none
    class(Records), intent(in) :: self
    type(Records) :: new_record

    allocate(new_record%num_infected_pupils, source=self%num_infected_pupils)
    
    allocate(new_record%num_in_class, source=self%num_in_class)
    
    allocate(new_record%infected_teacher, source=self%infected_teacher)
    
    allocate(new_record%susceptible, source=self%susceptible)
    
    allocate(new_record%exposed, source=self%exposed)
    
    allocate(new_record%unwell, source=self%unwell)
    
    allocate(new_record%quarantined, source=self%quarantined)
    
    allocate(new_record%recovered, source=self%recovered)
    
    if (allocated(self%contacts%time_step)) allocate(new_record%contacts%time_step, source=self%contacts%time_step)
    
    if (allocated(self%contacts%indxA)) allocate(new_record%contacts%indxA, source=self%contacts%indxA)
    
    if (allocated(self%contacts%indxB)) allocate(new_record%contacts%indxB, source=self%contacts%indxB)
    
    if (allocated(self%contacts%transmission)) allocate(new_record%contacts%transmission, source=self%contacts%transmission)
    
    allocate(new_record%cat_lists, source = self%cat_lists)

    allocate(new_record%secondary_infections, source = self%secondary_infections)

    allocate(new_record%pupil_absent_days, source = self%pupil_absent_days)

    new_record%total_infected_pupils = self%total_infected_pupils
    new_record%total_infected = self%total_infected

    new_record%num_classroom_transmission = self%num_classroom_transmission
    new_record%num_community_transmission = self%num_community_transmission
    return

  end function copyRecords

  subroutine destroy_Records(self)
    implicit none
    type(Records) :: self
    integer :: j

    ! write (*,*) 'destroy_Records called'

    if (allocated(self%num_infected_pupils)) deallocate(self%num_infected_pupils)
    if (allocated(self%num_in_class)) deallocate(self%num_in_class)
    if (allocated(self%infected_teacher)) deallocate(self%infected_teacher)
    if (allocated(self%susceptible)) deallocate(self%susceptible)
    if (allocated(self%exposed)) deallocate(self%exposed)
    if (allocated(self%unwell)) deallocate(self%unwell)
    if (allocated(self%quarantined)) deallocate(self%quarantined)
    if (allocated(self%recovered)) deallocate(self%recovered)

    if (allocated(self%contacts%time_step)) deallocate(self%contacts%time_step)
    if (allocated(self%contacts%indxA)) deallocate(self%contacts%indxA)
    if (allocated(self%contacts%indxB)) deallocate(self%contacts%indxB)
    if (allocated(self%contacts%transmission)) deallocate(self%contacts%transmission)

    if (allocated(self%cat_lists)) then
      do j=1,size(self%cat_lists)
        if (allocated(self%cat_lists(j)%susceptible)) deallocate(self%cat_lists(j)%susceptible)
        if (allocated(self%cat_lists(j)%exposed)) deallocate(self%cat_lists(j)%exposed)
        if (allocated(self%cat_lists(j)%unwell)) deallocate(self%cat_lists(j)%unwell)
        if (allocated(self%cat_lists(j)%quarantined)) deallocate(self%cat_lists(j)%quarantined)
        if (allocated(self%cat_lists(j)%recovered)) deallocate(self%cat_lists(j)%recovered)
        if (allocated(self%cat_lists(j)%infected)) deallocate(self%cat_lists(j)%infected)
        if (allocated(self%cat_lists(j)%inclass)) deallocate(self%cat_lists(j)%inclass)
      end do
    end if

    if (allocated(self%secondary_infections)) deallocate(self%secondary_infections)

    if (allocated(self%pupil_absent_days)) deallocate(self%pupil_absent_days)

    return
  end subroutine destroy_Records

end module classroom_records