Code cleanup and fixes