Updated with HW4 part 1