Code cleanup and finishing touches