some refactoring, fixed the bug definitely this time