Removing unused code and bug fixes - WIP