Some more logic, looks good