Small changes to logic