cleaning up code, bug fixes