Add doc string
Update class diagram