clean up code
post refactor with mary