Rework according to review