Some refactoring and bugfixes