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

      }

    }

  }

}