Cleaned up the code ... added some error checking