Cleanup code and add fix for battery boxes