added code removed during merge. Refactored a little.