Cleanup and refactoring

Small fixes and code cleanup.