Fixed #84

Fixed #84 and made the code a little simpler. Would be good to work out what to do with the repeated test on lines 777-778.