ENH: fixed according to the comments