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) + ")");

  }

}