Final refactoring... I hope