More changes due to updates refactor