Cleaned up the little code that we have