More refactoring work
Fixed some stuff