Added automatic integration time and tidied up the code a bit