#!/usr/bin/python3
import sys
import scipy.stats
import matplotlib.pyplot as plt
import numpy as np

# main function
if __name__ == "__main__":
    if len(sys.argv) != 5:
        error =  "usage: " + sys.argv[0] + " <k1> <k2> <n1> <n2>\n"
        error += "\tvaccinated group: k1 out of n1 become ill\n"
        error += "\tplacebo group:    k2 out of n2 become ill\n"
        sys.exit(error)

    k1 = int(sys.argv[1])
    k2 = int(sys.argv[2])
    n1 = int(sys.argv[3])
    n2 = int(sys.argv[4])
    a = 1.0 + float(n2)/n1

    if(k1<0 or k2<0 or n1<1 or n2<1):
        sys.exit("numbers should be non-negative: k1,k2>=0, n1,n2>=1")

    # in the limit of n1,n2 --> inf, one can show:
    # posterior distribution: P(efficacy<=x) = Beta[k2+1, k1+1](z)
    # with z = (a-1)/(a-x)  <==> x = a + (a-1)/z
    alpha = k2+1
    beta  = k1+1
    print(" estimated efficacy: %.1f%%" %
           (100.0*(1.0-float(k1)/n1/(float(k2)/n2))));
    print(" confidence intervals (limit case n1,n2-->inf):")
    for confidence in [0.95, 0.99, 0.999]:
        # two-sided confidence interval
        quantile1 = 0.5*(1.0-confidence);
        quantile2 = 1.0-quantile1;
        l = scipy.stats.beta.ppf(quantile1,alpha,beta)
        u = scipy.stats.beta.ppf(quantile2,alpha,beta)
        lower = a - (a-1.0)/l;
        upper = a - (a-1.0)/u;
        print("   %.2f%%: [%5.2f%%,%6.2f%%]" %
                (confidence*100.0, lower*100.0, upper*100.0));

    # plot
    x = np.linspace(0.0,1.0,500)
    plt.plot(x, scipy.stats.beta.pdf((a-1.0)/(a-x),alpha,beta) * (a-1.0)/((a-x)**2))
    plt.show()