prettified code, and added extra checks