Updated some logic to work better.