minor changes
Refactor: Replace conditional with polymorphism