Code refactor - final part. Bug fixing.