Clean up drawing code, add save/load functionality