Refactor code to get rid of conditional statements