Cleanup and refactor (currently broken)