Bug fixes and code cleanup