Code fix and improvement