Added if statement, fixed bugs