refactor the code, fix bugs