Fixed it to work the same as jonathanong's alternate version. I am
fairly convinced the algo is correct now.