Refactored code to fit checkstyle