Code cleanup and bug fixes