Some bugfixes and a few improvements

Let me know what you think!