Update for changes and cleanup