Refactor after Gilbert checked code