fix JS and clean up code