Refactor code