Changed some lines in JS to jQuery