refactor code for robustness
update code after review