Bug fix and small visual changes