Cleaned up the drawing code

By combining all the different models in 1 function