Update code to not suck