Improved some logic

Improved some logic and added some functionality