cleanup and bug fix part 2