Bug fix and a few tweaks