Code improvements with bug fixes