Fixed some errors and added some constraints to make the code more
robust