refactor if-else statment logic.

Just a little refactor. Have fun!