library(data.table)
library(dplyr)

indir = "~/git/R0t"

#
# Contact matrices from Zhang et al. Science https://science.sciencemag.org/content/368/6498/1481

pre_W = matrix(nrow = 14, ncol = 14, byrow = T,
               c(0.11, 0.35, 0.42, 0.55, 0.88, 0.35, 0.58, 0.12, 0.00, 0.00, 0.00, 0.00, 0.11, 0.33,
                 1.69, 2.31, 2.27, 5.84, 0.61, 0.58, 0.62, 0.55, 0.21, 0.23, 0.29, 0.17, 0.16, 0.18,
                 0.06, 3.76, 4.08, 8.39, 0.65, 0.75, 0.43, 0.31, 0.47, 0.60, 0.20, 0.31, 0.03, 0.16,
                 0.00, 1.35, 1.62, 4.41, 0.56, 0.43, 0.40, 0.24, 0.61, 0.47, 0.24, 0.28, 0.23, 0.14,
                 0.31, 0.38, 0.16, 0.82, 2.66, 1.65, 1.67, 1.37, 0.81, 0.98, 0.61, 0.67, 0.18, 0.41,
                 1.41, 0.70, 0.37, 0.99, 3.08, 1.67, 1.66, 1.43, 0.75, 0.97, 0.43, 0.48, 0.20, 0.45,
                 1.33, 1.29, 0.35, 1.11, 2.65, 1.82, 1.60, 1.26, 0.81, 1.15, 0.76, 0.62, 0.29, 0.48,
                 0.84, 0.99, 0.53, 1.36, 3.10, 1.92, 1.83, 1.70, 1.33, 1.45, 0.83, 0.83, 0.50, 0.86,
                 0.84, 0.60, 0.33, 0.83, 2.37, 1.20, 1.21, 1.23, 0.78, 1.30, 0.76, 0.58, 0.76, 1.60,
                 0.43, 0.88, 0.42, 1.39, 1.41, 0.76, 0.94, 0.91, 1.03, 1.18, 0.97, 0.79, 0.49, 0.88,
                 0.37, 0.58, 0.22, 0.75, 1.26, 0.94, 0.80, 0.60, 0.92, 1.20, 0.89, 0.66, 1.00, 2.03,
                 0.55, 0.69, 0.24, 1.07, 1.92, 0.98, 1.02, 1.34, 1.16, 0.62, 0.95, 0.84, 1.64, 4.10,
                 0.40, 0.24, 0.13, 0.39, 1.33, 0.73, 0.78, 0.61, 1.01, 1.76, 0.89, 0.76, 1.18, 2.76,
                 1.06, 0.68, 0.31, 0.86, 2.06, 1.50, 1.31, 0.69, 1.11, 1.17, 0.74, 1.22, 0.98, 1.91))

post_W = matrix(nrow = 14, ncol = 14, byrow = T,
                c(0.00, 0.11, 0.00, 0.00, 0.00, 0.00, 1.11, 0.44, 0.23, 0.00, 0.00, 0.00, 0.21, 0.11,
                  0.00, 0.00, 0.00, 0.00, 0.06, 0.05, 0.45, 0.87, 0.20, 0.10, 0.07, 0.09, 0.22, 0.09,
                  0.00, 0.00, 0.00, 0.04, 0.00, 0.00, 0.03, 0.78, 0.59, 0.32, 0.16, 0.03, 0.00, 0.07,
                  0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.43, 0.47, 0.72, 0.15, 0.00, 0.05, 0.20,
                  0.04, 0.02, 0.01, 0.05, 0.13, 0.05, 0.06, 0.02, 0.19, 0.43, 0.42, 0.21, 0.11, 0.10,
                  0.24, 0.16, 0.03, 0.04, 0.10, 0.31, 0.25, 0.10, 0.10, 0.04, 0.36, 0.21, 0.10, 0.03,
                  0.29, 0.38, 0.16, 0.01, 0.00, 0.07, 0.39, 0.24, 0.10, 0.02, 0.06, 0.11, 0.28, 0.19,
                  0.14, 0.30, 0.29, 0.06, 0.00, 0.02, 0.22, 0.31, 0.28, 0.04, 0.03, 0.04, 0.17, 0.19,
                  0.05, 0.07, 0.32, 0.27, 0.08, 0.01, 0.02, 0.11, 0.40, 0.27, 0.07, 0.02, 0.05, 0.30,
                  0.00, 0.08, 0.03, 0.24, 0.27, 0.16, 0.02, 0.08, 0.17, 0.37, 0.28, 0.09, 0.02, 0.41,
                  0.04, 0.01, 0.03, 0.03, 0.24, 0.23, 0.09, 0.05, 0.03, 0.13, 0.32, 0.25, 0.05, 0.21,
                  0.26, 0.00, 0.05, 0.00, 0.11, 0.20, 0.16, 0.05, 0.20, 0.00, 0.15, 0.31, 0.10, 0.36,
                  0.06, 0.10, 0.01, 0.01, 0.05, 0.01, 0.08, 0.08, 0.02, 0.06, 0.08, 0.14, 0.53, 0.23,
                  0.00, 0.06, 0.06, 0.00, 0.00, 0.00, 0.12, 0.12, 0.06, 0.06, 0.06, 0.06, 0.23, 0.47))


pre_S = matrix(nrow = 14, ncol = 14, byrow = T,
               c(3.17, 0.24,  0.06, 0.03,  0.07, 0.82, 1.21, 0.68, 0.17, 0.11, 0.26, 0.53, 0.63, 0.62,
                 0.29, 13.03, 1.76, 0.04,  0.13, 0.36, 0.93, 1.41, 0.26, 0.16, 0.04, 0.30, 0.57, 0.86,
                 0.00, 1.14,  17.5, 2.51,  0.25, 0.29, 0.47, 1.04, 0.92, 0.49, 0.23, 0.13, 0.16, 0.69,
                 0.00, 0.05,  0.97, 17.64, 0.85, 1.03, 1.64, 2.57, 1.65, 1.59, 0.94, 0.71, 0.28, 0.92,
                 0.10, 0.00,  0.08, 0.07,  2.15, 2.67, 2.44, 3.04, 2.95, 3.17, 1.72, 1.27, 0.61, 0.82,
                 0.16, 0.07,  0.07, 0.21,  1.14, 2.76, 2.56, 3.29, 2.05, 2.41, 1.86, 1.74, 0.87, 1.27,
                 0.42, 0.41,  0.06, 0.32,  1.44, 2.65, 4.11, 4.53, 2.81, 3.38, 1.98, 1.86, 1.05, 1.17,
                 0.11, 0.29,  0.33, 0.18,  0.76, 1.57, 2.30, 3.07, 2.17, 2.32, 1.31, 1.33, 0.69, 0.63,
                 0.04, 0.08,  0.37, 0.45,  0.76, 1.55, 2.20, 3.23, 2.21, 2.82, 1.34, 0.90, 0.37, 1.51,
                 0.04, 0.09,  0.25, 0.27,  1.42, 1.95, 2.64, 2.81, 2.78, 2.91, 2.03, 1.18, 0.58, 1.53,
                 0.09, 0.06,  0.12, 0.24,  1.04, 1.77, 1.48, 2.70, 2.15, 2.57, 2.10, 1.68, 1.48, 2.56,
                 0.10, 0.09,  0.04, 0.07,  0.61, 1.24, 1.66, 1.96, 1.32, 1.75, 1.46, 2.37, 1.08, 2.79,
                 0.05, 0.03,  0.10, 0.05,  0.26, 0.72, 0.96, 1.10, 0.82, 0.93, 1.03, 1.57, 1.78, 2.46,
                 0.03, 0.12,  0.06, 0.06,  0.13, 0.27, 0.25, 0.56, 0.69, 0.74, 0.68, 0.88, 1.48, 5.56))

post_S = matrix(nrow = 14, ncol = 14, byrow = T,
                c(0.00, 0.00, 0.00, 0.00, 0.00, 0.29, 0.55, 1.01, 0.14, 0.00, 0.14, 0.00, 0.00, 0.00,
                  0.04, 0.00, 0.00, 0.00, 0.00, 0.26, 0.51, 0.75, 0.11, 0.08, 0.00, 0.00, 0.00, 0.00,
                  0.00, 0.00, 0.00, 0.00, 0.14, 0.05, 0.38, 0.47, 0.34, 0.54, 0.09, 0.00, 0.00, 0.09,
                  0.12, 0.06, 0.19, 0.18, 0.25, 0.43, 0.11, 0.20, 0.72, 0.52, 0.43, 0.23, 0.09, 0.46,
                  0.02, 0.02, 0.06, 0.04, 0.08, 0.09, 0.05, 0.09, 0.25, 0.32, 0.42, 0.26, 0.13, 0.18,
                  0.05, 0.06, 0.00, 0.03, 0.09, 0.22, 0.11, 0.08, 0.08, 0.07, 0.35, 0.27, 0.04, 0.06,
                  0.18, 0.20, 0.05, 0.01, 0.08, 0.12, 0.57, 0.28, 0.16, 0.07, 0.10, 0.17, 0.25, 0.19,
                  0.15, 0.26, 0.19, 0.08, 0.05, 0.10, 0.15, 0.43, 0.35, 0.11, 0.04, 0.03, 0.31, 0.28,
                  0.16, 0.23, 0.31, 0.24, 0.11, 0.08, 0.15, 0.23, 0.54, 0.34, 0.14, 0.09, 0.21, 0.66,
                  0.00, 0.00, 0.28, 0.29, 0.23, 0.12, 0.11, 0.17, 0.25, 0.39, 0.33, 0.14, 0.00, 0.10,
                  0.02, 0.03, 0.07, 0.05, 0.24, 0.36, 0.07, 0.06, 0.00, 0.10, 0.43, 0.19, 0.22, 0.12,
                  0.03, 0.04, 0.00, 0.00, 0.16, 0.08, 0.22, 0.00, 0.23, 0.08, 0.16, 0.26, 0.35, 0.15,
                  0.02, 0.06, 0.02, 0.03, 0.06, 0.08, 0.21, 0.12, 0.13, 0.00, 0.02, 0.10, 0.40, 0.38,
                  0.00, 0.04, 0.01, 0.05, 0.00, 0.05, 0.10, 0.09, 0.06, 0.00, 0.00, 0.05, 0.28, 0.58))


# Define age categories

age_category = c("0-4", "5-9", "10-14", "15-19", "20-24", "25-29", "30-34", "35-39", "40-44", "45-49", "50-54", "55-59", "60-64", "65+")


# Save 

# Shanghai
tmp = data.table(age_index = rep(age_category, length(age_category)),
                 age_contact = rep(age_category, each = length(age_category) ),
                 count_pre = as.vector(pre_S),
                 count_post = as.vector(post_S) )
tmp[, city := "Shanghai"]

# Wuhan
tmp1 = data.table(age_index = rep(age_category, length(age_category)),
                  age_contact = rep(age_category, each = length(age_category) ),
                  count_pre = as.vector(pre_W),
                  count_post = as.vector(post_W) )
tmp1[, city := "Wuhan"]
tmp = rbind(tmp, tmp1)

tmp[, age_index := factor(age_index, levels = age_category)]
tmp[, age_contact := factor(age_contact, levels = age_category)]

saveRDS(tmp, file = file.path(indir, "usa", "data", "meta", "estimate_contact_intensities_outbreak_Zhang.rds"))


#
# Contact matrices from Jarvis et al. BMC https://bmcmedicine.biomedcentral.com/articles/10.1186/s12916-020-01597-8

library('socialmixr')

pre_UK = contact_matrix(polymod, countries = "United Kingdom", age.limits = c(0, 5, 18, 30, 40, 50, 60, 70))$matrix

post_UK = matrix(ncol = 8, nrow = 8, byrow = T,
                 c(0.525697998, 0.402583748, 0.368722035, 0.538440709, 0.194780472, 0.173757609, 0.098934288,	0.047254293,
                   0.14245221,	 2.430861989, 0.42609208	, 0.514546595, 0.48340504	, 0.188256272, 0.12084115	, 0.111607242,
                   0.137760561, 0.449900339, 1.194867894, 0.483464468, 0.480124784,	0.346225317, 0.184960293,	0.121814092,
                   0.209770231, 0.566523235, 0.504132528, 0.757773889, 0.52903225	, 0.346035983, 0.239739093,	0.08818435,
                   0.077825866, 0.545854342, 0.51346029	, 0.542568686, 0.795064526, 0.362278814, 0.233184095,	0.237601009,
                   0.079827148, 0.244423681, 0.425735426, 0.408058164, 0.416553864, 0.509931872, 0.264286413,	0.204649283,
                   0.058580903, 0.202213955, 0.293131116, 0.364369583, 0.345564886, 0.340625601, 0.321611885,	0.271976601,
                   0.023550101, 0.157192176, 0.162488722, 0.11280726 , 0.296361125, 0.22200101,	 0.228914771,	0.369024457))

# Define age categories

age_category = c("0-4", "5-17", "18-29", "30-39", "40-49", "50-59", "60-69", "70+")

tmp = data.table(age_index = rep(age_category, length(age_category)),
                 age_contact = rep(age_category, each = length(age_category) ),
                 count_pre = as.vector(pre_UK),
                 count_post = as.vector(post_UK) )

tmp[, age_index := factor(age_index, levels = age_category)]
tmp[, age_contact := factor(age_contact, levels = age_category)]

saveRDS(tmp, file = file.path(indir, "usa", "data", "meta", "estimate_contact_intensities_outbreak_UK.rds"))

