Cleaning up for new  features to be added