using Turing, StatsPlots, Distributions, Printf

@model function antibody_and_pcr_tests(
            results,n,deaths,population
            )
    # Model for Streeck et al
    #Prior for the ifr
    ifr ~ Uniform(0.0,0.05)
    # The population is grouped into 4 partitions
    # α - people who have antibodies but are not currently infected
    # β - people who have antibodies and are currently infected
    # γ - people who are currently infected but do not have antibodies
    # δ - people who are neither infected, nor have antibodies
    #Uniform priors
    subpopulations ~ Dirichlet([1,1,1,1])
    α = subpopulations[1]
    β = subpopulations[2]
    γ = subpopulations[3]
    δ = subpopulations[4]
    #
    #How many people test positive for antibody test from subpopulations α,β,γ,δ = (1,2,3,4) respectively
    #There is some uncertainty associated to these variables, which we attempt to model
    #These first variables are sensitivities
    # Streeck et al claim that sensitivity is 90.9%, but no number of samples given
    # So we assume it's something like 200 and of these 180 are positive and 20 negative
    antibody_test_1 ~ Beta(1+180,1+20)
    #For people that are currently infected but have antibodies we assume they have less antibodies, so smaller positivity rate.
    #Probably also more uncertainty with these numbers
    antibody_test_2 ~ Beta(1+80,1+20)
    #These people are infected but without substantial antibodies. Let's assume the false positive rate is slightly higher here
    #Than completely healthy people
    antibody_test_3 ~ Beta(1+20,1+80)
    #For these, Streeck et. all claim that 1656 samples of the manufacturer had 15 false positives
    #Then they themselves get 2 false positives for 109 tests
    #So in total, we know of 17 false positives and  1748 negatives
    antibody_test_4 ~ Beta(1+17,1+1748)
    #
    #Same numbers for the PCR test
    #
    #We know that the PCR test tests quite long even after infection, so let's allow for around 5% "false positives" amongst those with antibodies
    #of people who were infected but are not longer infected
    pcr_test_1 ~ Beta(1+5,1+95)
    #Let's say PCR has around 70% sensitivity, based on 500 tests.
    pcr_test_2 ~ Beta(1+350,1+150)
    pcr_test_3 ~ Beta(1+350,1+150)
    #slightly less than 1% false positives amongst complete negatives due to e.g. contamination etc seems reasonable
    #But we're not so sure, so we act like we've seen 1 false positive test after doing 150 tests
    pcr_test_4 ~ Beta(1+1,1+149)
    #
    #
    #Probability of getting a positive antibody test
    prob_antibody = α*antibody_test_1 + β*antibody_test_2 + γ*antibody_test_3 + δ*antibody_test_4
    #Probability of getting a positive PCR test
    prob_pcr = α*pcr_test_1 + β*pcr_test_2 + γ*pcr_test_3 + δ*pcr_test_4
    #Probability that both are positive (we assume a certain statistical independence here that is probably not warranted)
    prob_both = α*antibody_test_1*pcr_test_1 + β*antibody_test_2*pcr_test_2 + γ*antibody_test_3*pcr_test_3 + δ*antibody_test_4*pcr_test_4
    #
    categoryprobs = [
                    prob_antibody - prob_both,
                    prob_both,
                    prob_pcr - prob_both,
                    1 - prob_pcr - prob_antibody + prob_both
                    ]
    results ~ Multinomial(n,categoryprobs)
    #
    total_infected = α + β + γ
    #
    deaths ~ Binomial(population,ifr*total_infected)
    #
end


nsamples = 919 #Total number of samples
#We ignore the two old PCR tests that weren't serology positive, as we do account for serology false-negatives already
results = [102,10,23, nsamples - 102 - 23 - 10]
#Total population and deaths in the region
population = 12597
deaths = 7
#
iterations = 500000
chain = sample(
    antibody_and_pcr_tests(
        results,nsamples,deaths,population
        ),
    HMC(0.03,10), iterations, progress=true);


total_positive_res = chain["subpopulations[1]"] .+ chain["subpopulations[2]"] .+ chain["subpopulations[3]"]

fig1 = density(total_positive_res,xlim=(0.05,0.2),label="",title="Infection rate posterior",color=:black)
q1t = quantile(total_positive_res[:],0.025)
q2t = quantile(total_positive_res[:],0.975)
vline!([q1t], label=(@sprintf "2.5%% quantile = %.2f%%" q1t*100),color=:red,legend=:topleft)
vline!([q2t], label=(@sprintf "97.5%% quantile = %.2f%%" q2t*100),color=:red,legend=:topleft)

ifr = (collect(chain[:ifr].data)[:])
fig2 = density(ifr,xlim=(0.0,0.01),label="",title="IFR posterior",color=:black)
q1 = quantile(ifr, 0.025)
q2 = quantile(ifr,0.975)
vline!([q1], label=(@sprintf "2.5%% quantile = %.2f%%" q1*100),color=:red,legend=:topright)
vline!([q2], label=(@sprintf "97.5%% quantile = %.2f%%" q2*100),color=:red,legend=:topright)

fig3 = plot(fig1,fig2)