Some code refractoring. Fixed a minor bug.