Revert "Clean up of class diagram"

This reverts commit 71afaadc268fdd4621f6396fba1080c822804d49.