Cleaned up code

Made it a bit easier to stage now.