Added more work on milestone 2