polishing code , small changes