Code clean up and bug fixes.