Cleanup of ugly code