cleaned up logic

added some functionality