Almost finished hw4 (needs some fixes)