added if statement and while loop