fixed bugs after refactoring :)