updating code with help from solutions branch