Minor changes, tested the code.