Code cleanup. Small bugs fixes.