Clean up and put in problem 6