Code rewrite

re-wrote the code to include a switch case instead of if statements