Major code review/clean up