Added functionality for second homework. Will refactor for decent code later