Add graphical stuff and code cleaning