Some partial refactor. Tests failing