Add some code from the demo