Changed logic flow a bit.