refactor and fixed bug