sclass ProbabilisticMachine2<A extends ProbabilisticMachine2.State> implements Steppable {
  transient TreeSetWithDuplicates<A> states = new(byProbability());
  transient TreeSetWithDuplicates<A> steppableStates = new(byProbability());
  transient TreeSetWithDuplicates<A> droppedStates = new(byProbability());
  transient int stateCount;

  bool verbose;
  double cutoffPercentage = 50;
  new L<IVF1<A>> onStateAdded;
  
  new ThreadLocal<A> currentState;
  new ThreadLocal<Double> currentProbability;
  
  Comparator<A> byProbability() { ret (a, b) -> cmp(b.probability, a.probability); }
  
  abstract sclass State {
    ProbabilisticMachine2 probabilisticMachine;
    int stateNumber;
    State previousState;
    double probability = 100;
    
    abstract State makeDescendant();

    // override me unless it's a terminal state
    run {}
    
    toString {
      ret toStringWithFields(this, "stateNumber", "probability");
    }
  
    State prepareDescendant(State s) {
      copyFields(this, s, 'probabilisticMachine, 'probability);
      s.previousState = this;
      ret s;
    }
    
    ProbabilisticMachine2 probabilisticMachine() {
      ret probabilisticMachine;
    }
    
    void switchTo(State state) {
      probabilisticMachine.currentState.set(state);
    }
    
    void addState(State state) {
      probabilisticMachine.addState(state);
    }
  }
  
  // This only works when A = ProbabilisticMachine2.State
  /*A addState(Runnable r) {
    if (r != null)
      ret addState((A) new State {
        void step { r.run(); }
      });
    null;
  }*/
  
  A addState(A s) {
    if (verbose) print("Adding state to machine " + this + ": " + s);
    s.probabilisticMachine = this;
    if (s.stateNumber == 0) s.stateNumber = ++stateCount;
    if (s.probability < cutoffPercentage) ret s with droppedStates.add(s);
    addToCollections(s, states, steppableStates);
    pcallFAll(onStateAdded, s);
    if (verbose) printStats();
    ret s;
  }
  
  public bool step() {
    ret stepFirstUnstepped();
  }
  
  // returns false when done stepping
  bool stepFirstUnstepped() {
    A s = popFirst(steppableStates), ret false if null;
    temp tempSetTL(currentState, s);
    temp tempSetTL(currentProbability, s.probability);
    onRunningState(s);
    try {
      s.run();
      true;
    } finally {
      onRanState(s);
    }
  }
  
  swappable void onRunningState(A s) {}
  swappable void onRanState(A s) {}

  void reset {
    clearAll(states, steppableStates, droppedStates);
    stateCount = 0;
  }

  void think {
    while ping (stepFirstUnstepped()) {}
  }

  void printStats() {
    print("States: " + stats(states) + ", steppable: " + stats(steppableStates) + ", dropped: " + stats(droppedStates));
  }
  
  S stats(TreeSetWithDuplicates<A> ts) {
    ret l(ts) + (empty(ts) ? "" : " (best: " + iround(first(ts).probability) + ")");
  }
}