edited code to fix code review issues