test_that("contactMatrixAgeSchool() works correctly for one school age group with two schools ", {
  agelims <- c(0, 5, 18)
  agepops <- c(100, 240, 1000)
  cmp <- contactMatrixPolymod(agelims, agepops)

  schoolagegroups <- c(2, 2)
  schoolpops <- c(180, 60)
  schportion <- 0.8

  schoolfrac <- schoolpops / sum(schoolpops)

  cmas <- contactMatrixAgeSchool(agelims, agepops, schoolagegroups, schoolpops, schportion)
  expect_equal(cmas[1, 1], cmp[1, 1])
  expect_equal(cmas[1, 2], cmp[1, 2] * schoolfrac[1])
  expect_equal(cmas[1, 3], cmp[1, 2] * schoolfrac[2])
  expect_equal(cmas[1, 4], cmp[1, 3])
  expect_equal(cmas[2, 1], cmp[2, 1])
  expect_equal(cmas[2, 2], cmp[2, 2] * (schportion + (1 - schportion) * schoolfrac[1]))
  expect_equal(cmas[2, 3], cmp[2, 2] * (1 - schportion) * schoolfrac[2])
  expect_equal(cmas[2, 4], cmp[2, 3])
  expect_equal(cmas[3, 1], cmp[2, 1])
  expect_equal(cmas[3, 2], cmp[2, 2] * (1 - schportion) * schoolfrac[1])
  expect_equal(cmas[3, 3], cmp[2, 2] * (schportion + (1 - schportion) * schoolfrac[2]))
  expect_equal(cmas[3, 4], cmp[2, 3])
  expect_equal(cmas[4, 1], cmp[3, 1])
  expect_equal(cmas[4, 2], cmp[3, 2] * schoolfrac[1])
  expect_equal(cmas[4, 3], cmp[3, 2] * schoolfrac[2])
  expect_equal(cmas[4, 4], cmp[3, 3])
})

test_that("contactMatrixAgeSchool() works correctly for one school age group with three schools
          and two pre-school age groups", {
  agelims <- c(0, 2, 5, 18)
  agepops <- c(50, 50, 240, 1000)
  cmp <- contactMatrixPolymod(agelims, agepops)

  schoolagegroups <- c(3, 3, 3)
  schoolpops <- c(120, 80, 40)
  schportion <- 0.8

  schoolfrac <- schoolpops / sum(schoolpops)

  cmas <- contactMatrixAgeSchool(agelims, agepops, schoolagegroups, schoolpops, schportion)
  expect_equal(cmas[1, 1], cmp[1, 1])
  expect_equal(cmas[1, 2], cmp[1, 2])
  expect_equal(cmas[1, 3], cmp[1, 3] * schoolfrac[1])
  expect_equal(cmas[1, 4], cmp[1, 3] * schoolfrac[2])
  expect_equal(cmas[1, 5], cmp[1, 3] * schoolfrac[3])
  expect_equal(cmas[1, 6], cmp[1, 4])
  expect_equal(cmas[2, 1], cmp[2, 1])
  expect_equal(cmas[2, 2], cmp[2, 2])
  expect_equal(cmas[2, 3], cmp[2, 3] * schoolfrac[1])
  expect_equal(cmas[2, 4], cmp[2, 3] * schoolfrac[2])
  expect_equal(cmas[2, 5], cmp[2, 3] * schoolfrac[3])
  expect_equal(cmas[2, 6], cmp[2, 4])
  expect_equal(cmas[3, 1], cmp[3, 1])
  expect_equal(cmas[3, 2], cmp[3, 2])
  expect_equal(cmas[3, 3], cmp[3, 3] * (schportion + (1 - schportion) * schoolfrac[1]))
  expect_equal(cmas[3, 4], cmp[3, 3] * (1 - schportion) * schoolfrac[2])
  expect_equal(cmas[3, 5], cmp[3, 3] * (1 - schportion) * schoolfrac[3])
  expect_equal(cmas[3, 6], cmp[3, 4])
  expect_equal(cmas[4, 1], cmp[3, 1])
  expect_equal(cmas[4, 2], cmp[3, 2])
  expect_equal(cmas[4, 3], cmp[3, 3] * (1 - schportion) * schoolfrac[1])
  expect_equal(cmas[4, 4], cmp[3, 3] * (schportion + (1 - schportion) * schoolfrac[2]))
  expect_equal(cmas[4, 5], cmp[3, 3] * (1 - schportion) * schoolfrac[3])
  expect_equal(cmas[4, 6], cmp[3, 4])
  expect_equal(cmas[5, 1], cmp[3, 1])
  expect_equal(cmas[5, 2], cmp[3, 2])
  expect_equal(cmas[5, 3], cmp[3, 3] * (1 - schportion) * schoolfrac[1])
  expect_equal(cmas[5, 4], cmp[3, 3] * (1 - schportion) * schoolfrac[2])
  expect_equal(cmas[5, 5], cmp[3, 3] * (schportion + (1 - schportion) * schoolfrac[3]))
  expect_equal(cmas[5, 6], cmp[3, 4])
  expect_equal(cmas[6, 1], cmp[4, 1])
  expect_equal(cmas[6, 2], cmp[4, 2])
  expect_equal(cmas[6, 3], cmp[4, 3] * schoolfrac[1])
  expect_equal(cmas[6, 4], cmp[4, 3] * schoolfrac[2])
  expect_equal(cmas[6, 5], cmp[4, 3] * schoolfrac[3])
  expect_equal(cmas[6, 6], cmp[4, 4])
})

test_that("contactMatrixAgeSchool() works correctly for two school age groups each with two schools ", {
  agelims <- c(0, 5, 13, 18)
  agepops <- c(100, 240, 200, 1000)
  cmp <- contactMatrixPolymod(agelims, agepops)

  schoolagegroups <- c(2, 2, 3, 3)
  schoolpops <- c(140, 100, 80, 120)
  schportion <- 0.77

  schoolfrac <- c(schoolpops[1:2] / sum(schoolpops[1:2]), schoolpops[3:4] / sum(schoolpops[3:4]))

  cmas <- contactMatrixAgeSchool(agelims, agepops, schoolagegroups, schoolpops, schportion)
  expect_equal(cmas[1, 1], cmp[1, 1])
  expect_equal(cmas[1, 2], cmp[1, 2] * schoolfrac[1])
  expect_equal(cmas[1, 3], cmp[1, 2] * schoolfrac[2])
  expect_equal(cmas[1, 4], cmp[1, 3] * schoolfrac[3])
  expect_equal(cmas[1, 5], cmp[1, 3] * schoolfrac[4])
  expect_equal(cmas[1, 6], cmp[1, 4])
  expect_equal(cmas[2, 1], cmp[2, 1])
  expect_equal(cmas[2, 2], cmp[2, 2] * (schportion + (1 - schportion) * schoolfrac[1]))
  expect_equal(cmas[2, 3], cmp[2, 2] * (1 - schportion) * schoolfrac[2])
  expect_equal(cmas[2, 4], cmp[2, 3] * schoolfrac[3])
  expect_equal(cmas[2, 5], cmp[2, 3] * schoolfrac[4])
  expect_equal(cmas[2, 6], cmp[2, 4])
  expect_equal(cmas[3, 1], cmp[2, 1])
  expect_equal(cmas[3, 2], cmp[2, 2] * (1 - schportion) * schoolfrac[1])
  expect_equal(cmas[3, 3], cmp[2, 2] * (schportion + (1 - schportion) * schoolfrac[2]))
  expect_equal(cmas[3, 4], cmp[2, 3] * schoolfrac[3])
  expect_equal(cmas[3, 5], cmp[2, 3] * schoolfrac[4])
  expect_equal(cmas[3, 6], cmp[2, 4])
  expect_equal(cmas[4, 1], cmp[3, 1])
  expect_equal(cmas[4, 2], cmp[3, 2] * schoolfrac[1])
  expect_equal(cmas[4, 3], cmp[3, 2] * schoolfrac[2])
  expect_equal(cmas[4, 4], cmp[3, 3] *(schportion + (1 - schportion) * schoolfrac[3]))
  expect_equal(cmas[4, 5], cmp[3, 3] * (1 - schportion) * schoolfrac[4])
  expect_equal(cmas[4, 6], cmp[3, 4])
  expect_equal(cmas[5, 1], cmp[3, 1])
  expect_equal(cmas[5, 2], cmp[3, 2] * schoolfrac[1])
  expect_equal(cmas[5, 3], cmp[3, 2] * schoolfrac[2])
  expect_equal(cmas[5, 4], cmp[3, 3] * (1 - schportion) * schoolfrac[3])
  expect_equal(cmas[5, 5], cmp[3, 3] * (schportion + (1 - schportion) * schoolfrac[4]))
  expect_equal(cmas[5, 6], cmp[3, 4])
  expect_equal(cmas[6, 1], cmp[4, 1])
  expect_equal(cmas[6, 2], cmp[4, 2] * schoolfrac[1])
  expect_equal(cmas[6, 3], cmp[4, 2] * schoolfrac[2])
  expect_equal(cmas[6, 4], cmp[4, 3] * schoolfrac[3])
  expect_equal(cmas[6, 5], cmp[4, 3] * schoolfrac[4])
  expect_equal(cmas[6, 6], cmp[4, 4])
})