Updates based on @alexander-yakushev comments