Cleaned up my code