// TODO: find splitPoint faster (just access the median element!)
// TODO: really slow making the initial tree (5s for ~250 points)

sclass PtTree extends AbstractSet<Pt> {
  int size;
  Node root = new Leaf(null);
  
  // config
  
  int maxPointsPerNode = 4;
    
  abstract class Node {
    Node parent;
    
    abstract bool add(Pt p);
    abstract void collectPointsIn(Rect r, L<Pt> out);
    
    abstract void replaceChild(Node from, Node to);
    
    void addAll(Iterable<Pt> l) {
      fOr (Pt p : l) add(p);
    }
  }
  
  class Leaf extends Node {
    Cl<Pt> points;
    
    *(Node *parent) {}
    
    bool add(Pt p) {
      if (main contains(points, p)) false;
      if (points == null) points = new L;
      points.add(p);
      ++size;
      possiblySplit();
      true;
    }
    
    void possiblySplit {
      if (l(points) > maxPointsPerNode)
        split();
    }
    
    void collectPointsIn(Rect r, L<Pt> out) {
      fOr (p : points)
        if (containsPt(r, p))
          out.add(p);
    }
    
    record noeq ProposedSplit(int dimension, int splitPoint, int count) {
      int error() { ret abs(count-half(l(points))); }
    }
    
    Split split() {
      int n = l(points), half = n/2;
      
      ProposedSplit xSplit = checkSplit(0);
      ProposedSplit ySplit = checkSplit(1);
      
      var best = xSplit.error() < ySplit.error() ? xSplit : ySplit;

      new Split split;
      split.dimension = (byte) best.dimension;
      split.splitPoint = best.splitPoint;
      Leaf a = new Leaf(split);
      IPred<Pt> pred = p -> ptCoord(split.dimension, p) >= split.splitPoint;
      a.addAll(antiFilter(pred, points));
      split.a = a;
      Leaf b = new Leaf(split);
      b.addAll(filter(pred, points));
      split.b = b;
      
      a.possiblySplit();
      b.possiblySplit();
      
      ret replaceMeWith(split);
    }
    
    <A extends Node> A replaceMeWith(A node) {
      node.parent = parent;
      if (parent != null) parent.replaceChild(this, node);
      else if (root == this) root = node;
      ret node;
    }

    ProposedSplit checkSplit(int dimension) {
      L<Pt> lx = sortedBy(points, p -> ptCoord(dimension, p));
      new ProposedSplit ps;
      ps.dimension = dimension;
      
      int n = l(points), half = n/2;
      int i = 0;
      int splitPoint = Int.MIN_VALUE;
      while true {
        int lastSplitPoint = splitPoint;
        splitPoint = ptCoord(lx.get(i++), dimension)+1;
        int lastI = i;
        while (i < n && ptCoord(lx.get(i), dimension) < splitPoint) ++i;
        if (i >= half) {
          if (abs(lastI-half) < abs(i-half)) {
            ps.splitPoint = lastSplitPoint;
            ps.count = lastI;
          } else {
            ps.splitPoint = splitPoint;
            ps.count = i;
          }
          ret ps;
        }
      }
    }
    
    void replaceChild(Node from, Node to) { unimplemented(); } // doesn't need
  }
  
  static final int minInf = Int.MIN_VALUE/4;
  static final int inf = Int.MAX_VALUE/4;
  
  class Split extends Node {
    byte dimension; // 0 for X, 1 for Y
    int splitPoint;
    Node a, b;
    
    bool add(Pt p) {
      ret (ptCoord(p, dimension) >= splitPoint ? b : a).add(p);
    }
    
    Rect aRect() {
      ret rectFromPoints(minInf, minInf,
        dimension == 0 ? splitPoint : inf,
        dimension == 1 ? splitPoint : inf);
    }
    
    Rect bRect() {
      ret rectFromPoints(
        dimension == 0 ? splitPoint : minInf,
        dimension == 1 ? splitPoint : minInf,
        inf, inf);
    }
    
    void collectPointsIn(Rect r, L<Pt> out) {
      Rect rA = intersectRect(r, aRect());
      if (nempty(rA))
        a.collectPointsIn(rA, out);
        
      Rect rB = intersectRect(r, bRect());
      if (nempty(rB))
        b.collectPointsIn(rB, out);
    }
    
    void replaceChild(Node from, Node to) {
      if (a == from) a = to;
      if (b == from) b = to;
    } 
  }
  
  public bool add(Pt p) {
    ret root.add(p);
  }
  
  L<Pt> pointsIn aka pointsInRect(Rect r) {
    new L<Pt> out;
    root.collectPointsIn(r, out);
    ret out;
  }
  
  public bool contains(O p) {
    ret p instanceof Pt && contains(p/Pt);
  }
  
  bool contains(Pt p) {
    ret nempty(pointsIn(rect(p.x, p.y, 1, 1)));
  }
  
  // use this to make a PtTree!
  // It's actually makes a balanced tree.
  // It is assumed that points contains no duplicates
  static PtTree fromPointSet(Iterable<Pt> points) {
    new PtTree tree;
    Leaf root = cast tree.root;
    root.points = cloneList(points);
    tree.size = l(root.points);
    root.possiblySplit();
    ret tree;
  }
  
  public int size() { ret size; }
  
  public Iterator<Pt> iterator() {
    /* What we actually want is a coroutine yielding the values.
    // Like this, if we had this syntax:
    
    ret coroutine {
      void iterate(Node node) {
        if (node cast Leaf)
          returnAll(node.points);
        else {
          cast node to Split;
          iterate(node.a);
          iterate(node.b);
        }
        
        run { iterate(root); }
      }
    };
    
    // So how do we turn this into an iterator?
    // Maybe through VStack.Computable.
    */
    
    new Var<Iterator<Pt>> out;
    new VStack stack;
    stack.push(new co_iterate(root, out));
    
    ret iff_endOnNull(new IF0<Pt> {
      bool done;
      
      public Pt get() {
        while ping (true) {
          // something in the pipe? return it
          if (out.has())
            if (out->hasNext())
              ret out->next();
            else
              out.clear(); // and clean the pipe afterwards, boy!
              
          if (done) null; // end of the line
          // nothing in the pipe, gotta walk the tree some more
          if (!step(stack))
            set done; // prepare to get off the train
        }
      }
    });
  }
  
  record noeq co_iterate(Node node, Var<Iterator<Pt>> out) extends VStackComputableWithStep {
    void step(VStack stack) {
      if (step == 0) {
        if (node cast Leaf) {
          out.set(main iterator(node.points));
          stack.return();
        } else {
          cast node to Split;
          stack.call(new co_iterate(node.a, out));
          ret with step = 1;
        }
      } else {
        cast node to Split;
        stack.replace(new co_iterate(node.b, out));
      }
    }
  }

}