Cleaned up code and added utility class