Cleaned the code even more