Fixing logic, take 2