// 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); }
}