Code clean up and bug fixing.