little changes in code, refactor