Bug fix and code cleanup