a little clean up and error fix