// Now probabilistic - i.e.

//   transformers can return WithProbability or call lowerProbability().

// B is the type of object we handle.

// Hopefully this class is persistent?

sclass TransformersOnObjects<B> implements Steppable {

  replace Trail with O.

  // using the oldschool type F1 here because it's a bit easier to handle

  new AllOnAll_probabilistic<F1, B> allOnAll;

  new Set<B> bs;

  new MultiSetMap<B, Trail> trails;

  new Map<O, IF2> backTransformers; // reverse of each transformer

  new Map<O, Double> probabilities;

  bool verboseAddObject = true;

  transient new ThreadLocal<Double> currentProbability;

  srecord TransformationTrail(O transformer, O argument) {

    F1 backTransformer; // really complex back transformer

    toString { ret squareBracket(str(transformer)) + " on " + squareBracket(str(argument)); }

  }

  // If this is true and a transformer returns an Iterable,

  // we unpack it and discard the Iterable.

  bool autoUnpackIterables = true;

  // make a trail for each transformer application

  bool autoTrails = true;

  // keep multiple trails per object

  bool allowMultipleTrails = true;

  // when calling think()

  long maxSteps = 1000;

  AllOnAllOperation<F1, B> op = new(allOnAll, new TheOp);

  // make it a class so it's persistable

  class TheOp implements IVF2<F1, B> {

    public void get(F1 a, B b) {

      if (!functionCallableOn_nonSynthetic(a, b)) ret;

      temp tempSetTL(currentProbability, getProbability(a)*getProbability(b));

      O out = callFOpt(a, b);

      O trail = autoTrails ? makeTrailForTransformerApplication(a, b) : null;

      if (out == null) ret;

      if (out instanceof WithProbability) {

        lowerCurrentProbability(((WithProbability) out).probability);

        out = ((WithProbability) out)!;

        if (out == null) ret;

      }

      if (autoUnpackIterables && out instanceof Iterable) // don't use this anymore

        for (O x : (Iterable) out) addObject((B) x, trail);

      else

        addObject((B) out, trail);

    }

  }

  class ProbFunc implements IF2<F1, B, Double> {

    public Double get(F1 a, B b) {

      ret getProbability(a) * getProbability(b);

    }

  }

  *() {

    allOnAll.probabilityForPair = new ProbFunc;

  }

  bool addAll(Iterable<B> l) {

    bool change;

    fOr (B b : l) if (add(b)) set change;

    ret change;

  }

  bool add(B b) { ret addObject(b); }

  bool addObject(B b) {

    if (!addIfNotNull(bs, b)) false;

    double prob = or(currentProbability!, 1.0);

    setProbability(b, prob);

    if (verboseAddObject) print(prob == 1

      ? "Added: " + sfu(b)

      : "Added with " + WithProbability(prob, sfu(b)));

    allOnAll.addB(b);

    true;

  }

  void addObject(B b, O trail) {

    if (b == null) ret;

    add(b);

    if (trail != null && (allowMultipleTrails || !trails.containsKey(b)))

      trails.put(b, trail);

  }

  void addTransformer(F1 f) {

    allOnAll.addA(f);

  }

  // simple backtransformer only taking the transformed object

  void addTransformer(F1 f, F1 backTransformer) {

    addTransformer(f, backTransformer == null ? null

      : new SimplifiedBackTransformer(backTransformer));

  }

  srecord SimplifiedBackTransformer(F1 backTransformer) implements IF2 {

    public O get(O in, O out) { ret callFIfActuallyCallable(backTransformer, out); }

    toString { ret str(backTransformer); }

  }

  // complex backtransformer taking the original and the transformed object

  <A, C> void addTransformer(F1<A, C> f, IF2<A, C, O> backTransformer) {

    addTransformer(f);

    mapPut(backTransformers, f, backTransformer);

  }

  void addTransformer(VF2 f) {

    addTransformer(func(O a) {

      addTransformer(func(O b) {

        if (callableOn_nonSynthetic(f, "get", a, b))

          f.get(a, b);

        null;

      });

      null;

    });

  }

  void addTransformer(F2 f) {

    addTransformer(func(O a) {

      addTransformer(func(O b) {

        callableOn_nonSynthetic(f, "get", a, b) ? f.get(a, b) : null

      });

      null;

    });

  }

  public bool step() { ret op.step(); }

  L<B> getObjects() { ret allOnAll.getBs(); }

  O makeTrailForTransformerApplication(O transformer, B argument) {

    ret new TransformationTrail(transformer, argument);

  }

  void printWithTrails() {

    pnlMap(getObjects(), b -> {

      Set _trails = trails.get(b);

      ret sfu(b) + (empty(_trails) ? ""

        : " [" + nTrails(_trails) + ": " + joinWithComma(_trails) + "]");

    });

  }

  // only returns the actual TransformationTrails

  Cl<TransformationTrail> getTrails(B b) {

    ret instancesOf TransformationTrail(trails.get(b));

  }

  Cl<TransformationTrail> getTrailsBy(B b, O transformer) {

    ret filter(getTrails(b), t -> eq(((TransformationTrail) t).transformer, transformer));

  }

  TransformationTrail getTrailBy(B b, O transformer) {

    ret firstThat(getTrails(b), t -> eq(((TransformationTrail) t).transformer, transformer));

  }

  IF2 getBackTransformer(O transformer) {

    ret backTransformers.get(transformer);

  }

  void printTransformers() {

    pnl(allOnAll.getAs());

  }

  double getProbability(O o) {

    ret or(probabilities.get(o), 1.0);

  }

  // Note: works only for transformers that are not rewritten by the addTransformer method

  // Works for all objects

  <A> A setProbability(A o, double p) {

    if (o == null) null;

    Double pOld = probabilities.get(p);

    if (pOld != null) {

      if (p == 1.0)

        probabilities.remove(o);

      else if (pOld.doubleValue() != p)

        probabilities.put(o, p);

    } else if (p != 1.0)

      probabilities.put(o, p);

    ret o;

  }

  // p = 0.0 to 1.0

  void lowerProbability(double p) { lowerCurrentProbability(p); }

  void lowerCurrentProbability(double p) {

    currentProbability.set(currentProbability! * p);

  }

  void think { stepAllWithStats(this, maxSteps); }

}