refactored the demo to have a cleaner look