Bug fix

Clean up and bug fixes.