Fix messed up logic