Fixed button bug - needs fixed to work better