Bug fix and improvement