Fix for issues and code refactoring

Fixed issues #108 & #109
Refactored Code.