Updated working code