Fixed some bugs

It works nicely now :D