Changed stupid logic to smart logic