Tidying up
Implemented parts of TODO