Clean up and work on the drawing feature