Fixing some bugs, not perfect yet