refactoring code to fix bug. don't know if it works.