Added some code for extra credit