#!/bin/env python

import matplotlib.pyplot as plt

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):
    '''
        Bayesian approach
    '''
    n_positive_tests = 1863.6     # number of positive tests in the last 14 days
    n_tests = 7366.1                # 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("probability of meeting a covid positive in groupsize of n")
    plt.xlabel("group size / n")
    plt.ylabel("probability")
    plt.grid(linestyle='--', linewidth='0.1')
    plt.show()

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

plot_data(data)