tidy up program
change to inherited classes