Changed up some logic