Fixes #4 and other improvements