Some refactoring + comments