#!/bin/env python

import matplotlib.pyplot as plt

def get_data():
    n_14_day_tests = 7366.1
    n_14_day_positive_tests =  (18.2 + 7.1) / 100 * n_14_day_tests
    return (n_14_day_tests, n_14_day_positive_tests)

def p_meet_positive(groupsize:int):
    '''
        Naive approach, only works if each individual gets tested.
    '''
    active_cases = 777.69 # 7 day incidence rate
    population_size = 100000 # sample size -> per 100k
    p_positive = active_cases / population_size

    return (1 - ( 1 - p_positive )**groupsize) * 100

def p_meet_positive_bayes(groupsize:int, n_positive_tests:float, n_tests:float):
    '''
        Bayesian approach
    '''
    #n_positive_tests = n_tests_pos  # number of positive tests in the last 14 days
    #n_tests = n_tests_tot           # total number of tests in the last 14 days
    n_population = 100000           # sample size -> per 100k for switzerland
    p_test_positive =  0.3          # probability of a positive person getting tested (0.5 is optimistic...)

    p_positive_test = n_positive_tests / n_tests
    p_test = n_tests / n_population
    p_positive = p_positive_test * p_test / p_test_positive

    return (1 - ( 1 - p_positive )**groupsize) * 100

def plot_data(data):
    plt.plot(data)
    plt.title("likelihood of having a covid positive in groupsize of n")
    plt.xlabel("group size / n")
    plt.ylabel("probability / %")
    plt.grid(linestyle='--', linewidth='0.1')
    plt.show()

if __name__ == "__main__":

    tests, pos = get_data()

    data = []
    for i in range(1, 70):
        data.append(p_meet_positive_bayes(i, tests, pos))
        print(f'The chance of meeting a positive in a group of {i} is: {p_meet_positive_bayes(i, tests, pos)}%')

    plot_data(data)