Added final working code