update and fix following the comments