Fixed some bugs after further testing