Fix bugs based on Tony's code