Improvement/code refactor

Code refactor
- Code improvements