Refactoring, cleanup, and bug fixes