Code cleanup + improvements