Fixed legend issue.
Small legend fix and removed print statements in some places.
Merge pull request #8 from austin1howard/kamil-working-unstable

bugfixes from @kamilm