Refactor code a little.
Merge pull request #8 from AI-comp/refactoring

Refactor code a little