module classroom_transmission

  use :: iso_c_binding
  use :: fgsl
  use :: utilities
  use :: prob
  use :: params_mod

  implicit none

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

  public :: simulate

  type, public :: R0_params
      real(kind=c_double) :: shape
      real(kind=c_double) :: scale
    end type R0_params

  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
  end type records

  type, public :: ClassMember
    integer(kind=c_int) :: cohort=0 ! 1:pupil, 2:teacher, 3:TA, 4:nonclass
    integer(kind=c_int) :: infected = 0
    integer(kind=c_int) :: days_infected = 0
    integer(kind=c_int) :: days_isolated = 0
    integer(kind=c_int) :: incubation_time = 0
    integer(kind=c_int) :: symptomatic = 0
    real(kind=c_double) :: R0 = 3.0_dp
    real(kind=c_double) :: mitigation_factor = 1.0_dp
    real(kind=c_double) :: reduced_R0 = 3.0_dp
    integer(kind=c_int) :: infectious_time = 0
    real(kind=c_double) :: poisson_rate = 1.0_dp
    logical :: in_class = .True.
    logical :: immune = .False.
  contains
    procedure, public, pass(self) :: describe => describe_ClassMember
  end type ClassMember

  interface ClassMember
    module procedure init_ClassMember_params
  end interface ClassMember

  type, public :: ClassRoom
      type(Parameters) :: Params
      type(ClassMember), dimension(:), allocatable :: members
      integer(kind=c_int) :: time_step
      integer(kind=c_int) :: day
      type(R0_params) :: teacher_R0_params
      type(R0_params) :: pupil_R0_params
      integer(kind=c_int) :: num_classroom_transmission=0
      integer(kind=c_int) :: num_community_transmission=0
    contains
      private
      procedure, pass(self) :: getCounts
      generic, public :: counts => getCounts
      procedure, pass(self) :: getInClass
      generic, public :: get_inclass => getInClass
      procedure, pass(self) :: getInfected
      generic, public :: get_infected => getInfected
      procedure, pass(self) :: getNotInfected
      generic, public :: get_not_infected => getNotInfected
      procedure, pass(self) :: getExposed
      generic, public :: get_exposed => getExposed
      procedure, pass(self) :: updateMembers
      generic, public :: update_members => updateMembers
      procedure, pass(self) :: updateInClass
      generic, public :: update_in_class => updateInClass
      procedure, pass(self) :: updateTeacher
      generic, public :: update_teacher => updateTeacher
      procedure, pass(self) :: updateCommunity
      generic, public :: update_community => updateCommunity
      procedure, pass(self) :: updateClassroom
      generic, public :: update => updateClassroom
      final :: ClassRoom_destruct
  end type ClassRoom

  interface ClassRoom
    module procedure init_ClassRoom_Params
  end interface ClassRoom

contains

  function init_ClassMember_params(Params, cohort, R0, infected, symptomatic, days_infected) result(self)
    implicit none
    type(ClassMember) :: self
    type(Parameters), intent(in) :: Params
    integer(kind=c_int), intent(in) :: cohort
    real(kind=c_double), intent(in) :: R0
    integer(kind=c_int), intent(in) :: infected
    integer(kind=c_int), intent(in) :: symptomatic
    integer(kind=c_int), intent(in), optional :: days_infected

    integer(kind=c_int) :: days_infected_ = 5

    self%cohort = cohort
    self%infected = infected
    self%symptomatic = symptomatic
    self%days_isolated = 0

    self%incubation_time = Params%incubation_time
    self%R0 = R0

    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

    self%reduced_R0 = self%mitigation_factor * self%R0

    self%infectious_time = Params%infectious_time

    if (present(days_infected)) days_infected_ = days_infected
    self%days_infected = days_infected_

    self%poisson_rate = self%reduced_R0/self%infectious_time/3.0_dp
    self%in_class = .TRUE.
    self%immune = .FALSE.

    return
  
  end function init_ClassMember_params

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

    write(*,*) 'Cohort = ', self%cohort
    write(*,*) 'infected = ', self%infected
    write(*,*) 'days_infected = ', self%days_infected
    write(*,*) 'incubation_time = ', self%incubation_time
    write(*,*) 'symptomatic = ', self%symptomatic
    write(*,*) 'R0 = ', self%R0
    write(*,*) 'mitigation_factor = ', self%mitigation_factor
    write(*,*) 'reduced_R0 = ', self%reduced_R0
    write(*,*) 'infectious_time = ', self%infectious_time
    write(*,*) 'poisson_rate = ', self%poisson_rate
    write(*,*) 'in_class = ', self%in_class
    write(*,*) '----------------------'
    return
  end subroutine describe_ClassMember


  function init_ClassRoom_Params(Params, &
                                 teacher_R0_params, &
                                 pupil_R0_params) result(self)
    implicit none
    type(ClassRoom) :: self
    type(Parameters), intent(in) :: Params
    type(R0_params), intent(in) :: teacher_R0_params
    type(R0_params), intent(in) :: pupil_R0_params

    integer(kind=c_int) :: j
    real(kind=c_double) :: R0
    integer(kind=c_int) :: infected
    integer(kind=c_int) :: symptomatic
    integer(kind=c_int) :: days_infected

    self%Params = Params

    if (allocated(self%members)) deallocate(self%members)
    allocate(self%members(Params%class_size+1))

    self%teacher_R0_params = teacher_R0_params
    self%pupil_R0_params = pupil_R0_params

    self%num_community_transmission = 0
    self%num_classroom_transmission = 0

    do j=1,Params%class_size

      ! Is the pupil infected?
      infected = int(pdists%bernoulli_int_rvs(Params%child_prevalence), kind=c_int)

      ! Get this pupil's R0 value
      R0 = pdists%loglogistic_rvs(shape=pupil_R0_params%shape, &
                    scale=pupil_R0_params%scale)

      ! Will this pupil be symptomatic?
      symptomatic = int(pdists%bernoulli_int_rvs(Params%P_symptomatic), kind=c_int)

      ! If infected, how long have they been infected?
      if (infected==1) then
        self%num_community_transmission = self%num_community_transmission + 1
        if (symptomatic==1) then
          days_infected = int(pdists%randint(0,Params%incubation_time), kind=c_int)
        else
          days_infected = int(pdists%randint(0,Params%infectious_time), kind=c_int)
        end if
      else
        days_infected = 0
      end if

      self%members(j) = ClassMember(Params, 1, R0, infected, symptomatic, days_infected)
    end do

    ! Now setup a teacher
    ! Is the teacher infected?
    infected = int(pdists%bernoulli_int_rvs(Params%adult_prevalence), kind=c_int)

    ! Get this teacher's R0 value
    R0 = pdists%loglogistic_rvs(shape=teacher_R0_params%shape, &
                  scale=teacher_R0_params%scale)

    ! Will the teacher be symptomatic?
    symptomatic = int(pdists%bernoulli_int_rvs(Params%P_symptomatic), kind=c_int)

    if (infected==1) then
      self%num_community_transmission = self%num_community_transmission + 1
      if (symptomatic==1) then
        days_infected = int(pdists%randint(0,Params%incubation_time), kind=c_int)
      else
        days_infected = int(pdists%randint(0,Params%infectious_time), kind=c_int)
      end if
    else
      days_infected = 0
    end if

    self%members(Params%class_size+1) = ClassMember(Params, 2, R0, infected, symptomatic, days_infected)
    self%time_step = 0
    self%day = 0
    return
  end function init_ClassRoom_Params

  subroutine ClassRoom_destruct(self)
    type(ClassRoom) :: self
    if (allocated(self%members)) deallocate(self%members)
  end subroutine ClassRoom_destruct

  subroutine getCounts(self, num_infected_pupils, infected_teacher, num_in_class)
    class(ClassRoom), intent(in) :: self
    integer(kind=c_int), intent(out) :: num_infected_pupils
    integer(kind=c_int), intent(out) :: infected_teacher
    integer(kind=c_int), intent(out) :: num_in_class

    integer(kind=c_int64_t), dimension(:), allocatable :: infected
    integer(kind=c_int64_t), dimension(:), allocatable :: in_class

    call self%get_infected(infected)
    infected_teacher = 0
    if (allocated(infected)) then
      if (in_array(infected,self%Params%class_size+1)) then
        infected_teacher = 1
        call remove(infected,self%Params%class_size+1)
      end if

      num_infected_pupils = size(infected)
    else
      num_infected_pupils = 0
    end if

    in_class = self%get_inclass()
    num_in_class = size(in_class)

    return

  end subroutine getCounts

  function getInClass(self) result(inclass)
    class(Classroom), intent(in) :: self
    integer(kind=c_int64_t), dimension(:), allocatable :: inclass

    integer(kind=c_int64_t) :: j

    do j=1,self%Params%class_size
      if (self%members(j)%in_class)  call push(inclass, j)
    end do
    return
  end function getInClass

  subroutine getInfected(self, infected)
    class(Classroom), intent(in) :: self
    integer(kind=c_int64_t), dimension(:), allocatable, intent(out) :: infected

    integer(kind=c_int64_t) :: j

    do j=1,self%Params%class_size+1
      if (self%members(j)%infected==1) then
        call push(infected, j)
      end if
    end do
    return
  end subroutine getInfected

  function getNotInfected(self) result(not_infected)
    class(Classroom), intent(in) :: self
    integer(kind=c_int64_t), dimension(:), allocatable :: not_infected

    integer(kind=c_int64_t) :: j

    do j=1,self%Params%class_size+1
      if ((self%members(j)%infected==0) .and. (self%members(j)%in_class)) then
        call push(not_infected, j)
      end if
    end do
    return
  end function getNotInfected

  subroutine getExposed(self, exposed)
    class(Classroom), intent(in) :: self
    integer(kind=c_int64_t), dimension(:), allocatable, intent(out) :: exposed

    integer(kind=c_int64_t) :: j

    do j=1,self%Params%class_size+1
      if ((self%members(j)%infected==0) .and. (self%members(j)%in_class) &
          .and. (.not.self%members(j)%immune)) then
          call push(exposed, j)
      end if
    end do
    return
  end subroutine getExposed

  subroutine updateDisease_pupils(self)
    class(Classroom), intent(inout) :: self

    integer :: j
    type(ClassMember) :: m

    do j=1,self%Params%class_size
      m = self%members(j)
      if (m%infected==1) then
        self%members(j)%days_infected = self%members(j)%days_infected+1
        if ((self%members(j)%days_infected > self%members(j)%incubation_time) &
            .and. (m%symptomatic==1)) then
          self%members(j)%in_class = .False.
          self%members(j)%days_isolated = self%members(j)%days_isolated + 1
        end if
        if (self%members(j)%days_infected > self%members(j)%infectious_time) then
          self%members(j)%immune = .True.
          self%members(j)%infected = 0
        end if
        if (self%members(j)%days_isolated > 14) then
          self%members(j)%in_class = .True.
        end if
      end if
    end do
    return
  end subroutine updateDisease_pupils

  subroutine updateInClass(self)
    class(Classroom), intent(inout) :: self

    integer(kind=c_int64_t), dimension(:), allocatable :: infected
    integer(kind=c_int64_t), dimension(:), allocatable :: exposed

    integer(kind=c_int64_t), dimension(:), allocatable :: num_new_infections
    integer(kind=c_int64_t), dimension(:), allocatable :: new_infections

    integer(kind=c_int64_t), dimension(:), allocatable :: interactions

    integer :: n_infected
    integer :: inft_indx

    integer(kind=c_int64_t) :: this_new_infections

    integer :: j

    call self%get_infected(infected)
    call self%get_exposed(exposed)

    if (allocated(exposed)) then
      if (allocated(infected)) then
        n_infected = size(infected)
        allocate(num_new_infections(n_infected))

        do j=1,n_infected
          inft_indx = int(infected(j))
          this_new_infections = pdists%poisson_rvs( &
                                  self%members(inft_indx)%poisson_rate)
          num_new_infections(j) = min(this_new_infections, &
                                      int(size(exposed), kind=c_int64_t))
        end do

        do j=1,n_infected
          if (num_new_infections(j)>0) then
            allocate(interactions(num_new_infections(j)))
            interactions = pdists%choose(exposed, num_new_infections(j))
            call push(new_infections, interactions)
            deallocate(interactions)
          end if
        end do

        if (allocated(new_infections)) then
          call remove_repeats(new_infections)

          self%num_classroom_transmission = self%num_classroom_transmission + size(new_infections)
          do j=1,size(new_infections)
            self%members(new_infections(j))%infected = 1
          end do
          deallocate(new_infections)
        end if
        deallocate(num_new_infections)
      end if
    end if

    return

  end subroutine updateInClass

  subroutine updateTeacher(self)
    implicit none
    class(ClassRoom), intent(inout) :: self

    type(ClassMember) :: teacher
    type(ClassMember) :: new_teacher
    real(kind=c_double) :: R0
    integer(kind=c_int) :: symptomatic
    integer(kind=c_int) :: days_infected

    integer(kind=c_int) :: infected

    teacher = self%members(self%Params%class_size+1)
    if (teacher%infected==1) then
      if (teacher%symptomatic==1) then
        if (teacher%days_infected > teacher%incubation_time) then ! Teacher showing symptoms, so remove from class

          ! Get new R0 for replacement teacher
          R0 = pdists%loglogistic_rvs(shape=self%teacher_R0_params%shape, &
                                            scale=self%teacher_R0_params%scale)
          ! Is new teacher infected
          infected = int(pdists%bernoulli_int_rvs(self%Params%adult_prevalence), kind=c_int)
          ! Will the new teacher be symptomatic or asymptomatic?
          symptomatic = int(pdists%bernoulli_int_rvs(self%Params%P_symptomatic), kind=c_int)
          ! If teacher is currently infected, how long have they been infected?
          if (infected==1) then
            self%num_community_transmission = self%num_community_transmission + 1
            if (symptomatic==1) then
              days_infected = int(pdists%randint(0,teacher%incubation_time), kind=c_int)
            else
              days_infected = int(pdists%randint(0,teacher%infectious_time), kind=c_int)
            end if
          end if
          new_teacher = ClassMember(self%Params, 2, R0, infected, symptomatic, days_infected)
          self%members(self%Params%class_size+1) = new_teacher
        end if
      end if
    end if

    return

  end subroutine updateTeacher

  subroutine updateCommunity(self)
    implicit none
    class(ClassRoom), intent(inout) :: self

    integer :: j
    integer(kind=c_int) :: infected

    do j=1,self%Params%class_size+1
      if (self%members(j)%infected /= 1) then
        infected = int(pdists%bernoulli_int_rvs(self%Params%community_incidence), kind=c_int)
        self%members(j)%infected = infected
        if (infected==1) self%num_community_transmission = self%num_community_transmission + 1
      end if
    end do
    return
  end subroutine updateCommunity

  subroutine updateClassroom(self)
    implicit none
    class(ClassRoom), intent(inout) :: self
    integer(kind=c_int) :: day
    self%time_step = self%time_step + 1
    day = self%day + 1
    self%day = mod(day, int(7, kind=c_int))
    call self%update_members()
    if ((self%day /= 5) .and. (self%day /= 6)) then
      call self%update_in_class()
    end if
    call self%update_teacher()
    call self%update_community()

    return
  end subroutine updateClassroom

  subroutine simulate(Params, stats)
    type(Parameters), intent(in) :: Params
    type(records), intent(out) :: stats

    type(ClassRoom) :: cs
    type(R0_params) :: pupil_R0_params
    type(R0_params) :: teacher_R0_params

    integer :: j

    allocate(stats%num_infected_pupils(Params%days+1))
    allocate(stats%num_in_class(Params%days+1))
    allocate(stats%infected_teacher(Params%days+1))

    pupil_R0_params%shape = 3.279022
    pupil_R0_params%scale = 2.723258

    teacher_R0_params%shape = 3.134950
    teacher_R0_params%scale = 3.263443

    cs = ClassRoom(Params, pupil_R0_params, teacher_R0_params)

    ! call cs%members(1)%describe()
    ! call cs%members(2)%describe()
    ! call cs%members(Params%class_size+1)%describe()

    call cs%counts(stats%num_infected_pupils(1), &
                      stats%infected_teacher(1), &
                      stats%num_in_class(1))

    do j=1,Params%days
      call cs%update()

      call cs%counts(stats%num_infected_pupils(j+1), &
                        stats%infected_teacher(j+1), &
                        stats%num_in_class(j+1))
    end do


  end subroutine simulate

end module classroom_transmission