Cleaning up code base