Fixed some logic and added some checks.