fix in the second part of the logic