#!/usr/bin/env python
"""Generate corresponding infectious disease models
"""
from __future__ import division
import argparse
import math

def findbeta(radius, nu, S0, I0, n=15):
    alpha = math.pi * radius**2
    N = S0 + I0
    if radius == 0:
        return 0
    elif nu == 0:
        return 2 / N - 2 * math.exp(N * math.log(-alpha + 1) / 2) / N
    beta = -math.log(1 - math.pi*radius**2)
    for _ in xrange(n):
        eta = N - nu * (math.log(S0) - math.log(nu / beta) + 1) / beta
        beta = (1 - math.exp(eta * math.log(-alpha + 1))) / eta
    return beta

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--days', default=100, type=int,
        help='number of days to simulate')
    parser.add_argument('--population', metavar='N', default=5000, type=int,
        help='population size')
    parser.add_argument('--infected', metavar='I0', default=5, type=int,
        help='initial number of infected')
    parser.add_argument('--radius', default=0.003, type=float,
        help='radius of infection')
    parser.add_argument('--nu', default=0.0, type=float,
        help='rate of recovery')
    args = parser.parse_args()

    beta = findbeta(args.radius, args.nu,
                    args.population - args.infected,
                    args.infected)

    print "beta =", beta
    print ("./dynamics2.py --population %d --infected %d --beta %g --nu %g " +
           "--days %d --plot --save d.json") % (
        args.population, args.infected, beta, args.nu, args.days)
    print ("./agent2.py --population %d --infected %d --radius %g --gamma %g " +
           "--days %d --plot --save a.json") % (
        args.population, args.infected, args.radius, args.nu, args.days)

if __name__ == '__main__':
    main()