General code review and corrections