Some changes
Refactoring to be cleaner