Added functionality. But needs to be cleaned!!