Updated code, fixed bugs