[Fix] Refactoring and comments