Code cleaning & fixing