- Implemented code suggested in pull request by ChrisMissal
- Added some safety checks
- Cleaned up some code