Tried to finish code... hope it works