Added extra code and fixed some bugs