Cleaned up code according to comments on pull request