Rearranged some logic