Added main() function, added check that input is correct