Tried fixing bug a little more