Fixed demo and refactored some code.