Bug fix and code improvements.