fixed some code - please confirm I did not break anything

#114