Cleanup of logic