Fixed up some of the code base