Code cleanup and fixed some more bugs