Added some functuonality