Cleaned up code even more