Moved some code to look better