Code changes and cleanup