Restructured the code a little