Refactor drawing code