Remove some unnecessary code and refactor
More minor changes and fixes