Fixed #5 issue. Code clean up.