Undo parts of previous commits and improve code