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