Fix null reference exception
Some improvements