Added some code for drawing