cleaning up the code part 2