Add code into add(), need to make it function