cleaned up code a bit based on pull request comments