More improvements and cleanup