Finish up and fix #334