Refactor code, various improvements