Changed some code around