Refactor plotting & other shite