Some refactoring

Code was very hard to read and now it's a bit better.