Update code

Update and edit code.
Merge pull request #11 from onimini/patch-10

Update code