Cleaned up my code a bit