Refactoring.
Should fix #306.
Merge pull request #312 from angathorion/refactor

Refactoring and small fix