Make graphics square
Simplify and improve drawing code