Fixed code to be more like example on wiki