Improving code. Related to #6