#!/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) != 3:  
        sys.exit("usage: " + sys.argv[0] + " <k> <n>")

    k = int(sys.argv[1])
    n = int(sys.argv[2])
    if(n<k or n<1 or k<0):
        sys.exit("n=" + str(n) +", k="+ str(k) +", but they need to be positive with n>=k")

    print("input:  %i out of %i say yes" % (k,n));
    print("task:   estimate probability p of saying yes")
    print("result:")
    print(" point estimate: p  = %.2f%%" % (float(k)/n*100.0))

    # posterior distribution is Beta(k+1, n-k+1)
    alpha = k+1
    beta  = n-k+1
    print(" confidence intervals:")
    for confidence in [0.95, 0.99, 0.999]:
        if(k==0):
            # one-sided confidence interval
            quantile1=0.0;
            quantile2=confidence;
        elif(k==n):
            # one-sided confidence interval
            quantile1=1.0-confidence;
            quantile2=1.0;
        else:
            # two-sided confidence interval
            quantile1 = 0.5*(1.0-confidence);
            quantile2 = 1.0-quantile1;
        lower = scipy.stats.beta.ppf(quantile1,alpha,beta)
        upper = scipy.stats.beta.ppf(quantile2,alpha,beta)
        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(x,alpha,beta))
    plt.show()