Fix bugs and code review