Wrote logic for assignment 2. Just need to test it now.