Add back code i accidently removed while tinkering.