Added some graph logic