Code cleanup
Fixed a display problem