More cleaning up code