Assignment 1 code nearly done