code improvement and bug fixes