Fixed few more bugs