Fixes plus new features
tidy up of new feature
Merge pull request #15 from L1ghtn1ng/fixes

Fixes and feature
Clean up.