Cleanup, working code