Bug fix and various code improvements

Change-Id: I61e4d40dcefc80f4cff267e786de924eb20742c2