Final cleaned up code for Hw3