Fixed some issues. WIP