Fix/polish

Some comments and cosmetic changes