// As long as you don't access the by-numeral-index functions,

// you can add and remove elements quickly.

// Upon accessing the list, it is updated once to fit the new data.

// Note: When adding the same element twice, it will be stored only

//       once with the higher probabilities of the two.

//

// class is synchronized

sclass ProbabilisticList<A> extends AbstractRandomAccessList<WithProbability<A>> {

  // these are the 2 main data structures we update

  Map<A, Double> probabilities = hashMap();

  MultiSetMap<Double, A> byProbability

    = multiSetMap_outerDescTreeMap_innerCompactLinkedHashSet();

    // = multiSetMap_outerDescTreeMap_innerBetterLinkedHashSet();

  L<WithProbability<A>> list; // this is only updated when requested

  // CUSTOMIZATION

  // must be >= 0. probability 0 is never added

  double cutoffProbabilityOnAdd = 0;

  Cl<ProbabilisticListListener<A>> listeners;

  *() {}

  *(Iterable<WithProbability<A>> l) {

    main addAll(this, l);

  }

  public synchronized WithProbability<A> get(int idx) {

    if (idx < 0 || idx >= size()) null;

    makeListUpToIndex(idx);

    ret list.get(idx);

  }

  public synchronized int size() { ret l(probabilities); }

  // assuming the index is within range

  void makeListUpToIndex(int idx) {

    if (idx < l(list)) ret; // index already there

    if (list == null) list = new L;

    WithProbability<A> lastEntry = main last(list);

    A a = lastEntry?!;

    double probability = lastEntry == null ? 0 : lastEntry.probability();

    //print("Making list from " + l(list) + " to " + idx);

    while (idx >= l(list)) {

      //printVars_str(+idx, +a, +probability);

      if (a == null) {

        probability = byProbability.firstKey();

        a = first(byProbability.get(probability));

      } else {

        CompactLinkedHashSet<A> setForProbability = cast byProbability.get(probability);

        a = setForProbability.nextElement(a);

        //printVars_str(+probability, +setForProbability, +a);

        if (a == null) {

          probability = byProbability.higherKey(probability);

          a = first(byProbability.get(probability));

        }

      }

      list.add(WithProbability(probability, a));

    }

  }

  synchronized bool containsElement(A a) {

    ret probabilities.containsKey(a);

  }

  public bool add aka put(A a, double probability) {

    ret add(WithProbability(probability, a));

  }

  public bool add(WithProbability<A> wp) {

    A a;

    double newP;

    synchronized {

      if (wp == null || wp.probability() <= cutoffProbabilityOnAdd) false;

      a = wp!;

      newP = wp.probability();

      Double oldP = probabilities.get(a);

      if (oldP != null) {

        if (oldP.doubleValue() >= newP)

          false;

        else

          remove(WithProbability<A>(oldP, a));

      }

      clearListFrom(newP);

      probabilities.put(a, newP);

      byProbability.add(newP, a);

    }

    fireElementAdded(a, newP);

    true;

  }

  // will also delete if given probability doesn't match

  @Override

  public bool remove(O wp) {

    if (wp == null) false;

    cast wp to WithProbability<A>;

    ret removeElement(wp!);

  }

  public synchronized bool removeElement(A a) {

    if (a == null) false;

    Double oldP = probabilities.get(a);

    if (oldP == null) false;

    clearListFrom(oldP);

    probabilities.remove(a);

    byProbability.remove(oldP, a);

    true;

  }

  @Override

  public synchronized WithProbability<A> remove(int idx) {

    var a = get(idx);

    remove(a);

    ret a;

  }

  // internal

  void clearListFrom(double p) {

    int idx = generalizedBinarySearch2(list, a -> a.probability() <= p ? 1 : -1);

    idx = -idx-1;

    //printVars clearListFrom(p, idx, before := takeFirst(idx, list), after := main subList(idx, list));

    truncateList(list, idx);

  }

  void at(double p, A a) {

    add(WithProbability(p, a));

  }

  // faster version of truncateList

  void truncate(int size) {

    if (size < 0) fail();

    while (size() > size)

      removeLast();

  }

  double getProbability(A a) {

    ret orZero(probabilities.get(a));

  }

  void truncateBelow(double thresholdProbability) {

    A a;

    while ((a = last()) != null && getProbability(a) < thresholdProbability)

      removeElement(a);

  }

  A last() {  

    Double last = byProbability.lastKey();

    if (last == null) null;

    CompactLinkedHashSet<A> set = cast byProbability.getOpt(last);

    ret set.last();

  }

  void removeLast {

    removeElement(last());

  }

  int indexOfElement(A a) {

    // quickly bail if element is not contained

    double p = getProbability(a);

    if (p == 0) ret -1;

    // TODO: binary search and stuff

    int n = size();

    for i to n:

      if (eq(a, get(i)!))

        ret i;

    fail("Internal error");

  }

  // elements without probabilities

  L<A> rawList aka elementsOnly() {

    ret lazyMap_noSave(this, a -> a!);

  }

  // just the probabilities in the list

  L<Double> probabilities() {

    ret lazyMap_noSave(this, a -> a.probability());

  }

  L<ProbabilisticListListener<A>> listenersList() {

    if (listeners == null) null;

    synchronized { ret cloneList(listeners); }

  }

  void fireElementAdded(A a, double probability) {

    main forEach(listenersList(), l -> l.elementAdded(a, probability));

  }

  synchronized void addListener(ProbabilisticListListener<A> l) {

    listeners = addOrCreate(listeners, l);

  }

  synchronized void removeListener(ProbabilisticListListener<A> l) {

    listeners?.remove(l);

  }

  void onElementAdded(IVF2<A, Double> f) {

    if (f != null)

      addListener(new ProbabilisticListAdapter<A> {

        @Override

        public void elementAdded(A a, double probability) {

          f.get(a, probability);

        }

      });

  }

}