Fix review comments, improve formatting and styles