// Flood-filling everywhere with neighboring pixel similarity
// threshold instead of posterizing

// TODO: handle withDiagonals properly (do diagonal similarity testing)

abstract srecord noeq DifferentialRegionsMaker<Img extends WidthAndHeight>(Img image) is Runnable, IImageRegions<Img> {
  int w, h, runner;
  gettable int size; // =w*h
  
  new IntBuffer stack; // locations as y*w+x
  gettable int[] regionMatrix; // for each pixel: region index (starting at 1)
  new IntBuffer regionPixels; // collect all pixels for regions
  
  settable bool withDiagonals; // also walk diagonally?
  
  // How different pixels may be from their neighbor in "the
  // "same region. From 0 to 1.
  settable double tolerance = 0.05;
  
  // initialize these to use them
  
  new IntBuffer regionFirstPixel; // for each region: index of first pixel found in regionPixels
  new IntBuffer regionSize; // for each region: number of pixels
  new IntBuffer regionBounds; // for each region: bounds (x1, y1, x2, y2)
  
  // index = dual log of region size, value = region index
  new L<IntBuffer> regionsBySize;
  
  int regionCounter;
  bool verbose;
  
  double regionStep = .1; // for rendering in regionsImage
  
  int x(int pos) { ret pos % w; }
  int y(int pos) { ret pos / w; }
  int pos(int x, int y) { ret y*w+x; }
  Pt pt(int pos) { ret Pt(x(pos), y(pos)); }
  bool validPos(int x, int y) { ret x >= 0 && y >= 0 && x < w && y < h; }
  
  abstract int getRGB(int pos);
  
  bool similarPixels(int pos1, int pos2) {
    int col1 = getRGB(pos1);
    int col2 = getRGB(pos2);
    ret rgbDiff(col1, col2) <= tolerance;
  }
  
  run {
    if (regionMatrix != null) ret;
    
    if (withDiagonals) fail("Can't handle diagonals yet");
    
    w = image.getWidth(); h = image.getHeight();
    size = w*h;
    regionMatrix = new int[size];
    
    // 0 entries are unused
    regionFirstPixel?.add(0); 
    regionSize?.add(0);
    
    regionPixels?.setSize(size);
    
    while (runner < size) {
      if (regionMatrix[runner] != 0)
        continue with ++runner;
      
      makeRegion_fast();
    }
  }
  
  void makeRegion_fast {
    // make a new region
    int region = ++regionCounter;
    regionFirstPixel?.add(regionPixels != null ? l(regionPixels) : runner);
    stack.add(runner);
    int rsize = 0, x1 = w, y1 = h, x2 = 0, y2 = 0;
    ifdef FastRegions_debug
      printVars FastRegions(+region, +runner);
    endifdef

    // flood-fill region
    while (nempty(stack)) {
      int pos = stack.popLast();
      if (regionMatrix[pos] != 0) continue; // check again if set in the meantime (color has been checked before)

      // new pixel found!
      int x = x(pos), y = y(pos);
      
      // march to the left
      int lineStart = pos-x;
      int xLeft = x;
      while (xLeft > 0 && addable(lineStart+xLeft-1, -1))
        --xLeft;
        
      // march to the right
      int xRight = x+1;
      while (xRight < w && addable(lineStart+xRight, 1))
        ++xRight;
      
      // Now we know [xLeft; xRight) are our pixels
      // mark them in matrix
      for (x = xLeft; x < xRight; x++) {
        regionMatrix[lineStart+x] = region;
        regionPixels?.add(lineStart+x);
      }
        
      // increase region size
      rsize += xRight-xLeft;
      
      ifdef FastRegions_debug
        printVars FastRegions(+region, +rsize, +y, +xLeft, +xRight);
      endifdef
      
      // update bounds
      if (xLeft < x1) x1 = xLeft;
      if (xRight-1 > x2) x2 = xRight-1;
      if (y < y1) y1 = y;
      if (y > y2) y2 = y;
      
      // explore above & below
      int xLeft2 = withDiagonals ? max(0, xLeft-1) : xLeft;
      int xRight2 = withDiagonals ? min(w, xRight+1) : xRight;
      if (y > 0)
        addStreaks(lineStart-w, xLeft2, xRight2, -w);
      if (y < h-1)
        addStreaks(lineStart+w, xLeft2, xRight2, w);
    }
    
    regionSize?.add(rsize);
    regionBounds?.addAll(x1, y1, x2+1, y2+1);
    if (regionsBySize != null) {
      int iBucket = dualLog(rsize);
      var buffer = listGetOrCreate(regionsBySize, iBucket, -> new IntBuffer);
      buffer.add(region);
    }
  }
  
  bool addable(int pos, int lastDirection) {
    if (regionMatrix[pos] != 0) false; // touching myself (or someone else)
    // check color similarity
    if (!similarPixels(pos, pos-lastDirection)) false;
    true;
  }
  
  bool addIfAddable(int pos, int lastDirection) {
    if (!addable(pos, lastDirection)) false;
    stack.add(pos);
    true;
  }
  
  void addStreaks(int lineStart, int xLeft, int xRight, int lastDirection) {
    int x = xLeft;
    while (x < xRight) {
      if (addIfAddable(lineStart+x, lastDirection))
        while (x+1 < xRight
          && addable(lineStart+x+1, 1)
          && addable(lineStart+x+1, lastDirection))
          ++x;
      ++x;
    }
  }

  IBWImage regionsImage() {
    ret iBWImageFromFunction((x, y) -> {
      var region = regionMatrix[pos(x, y)];
      ret ((region-1)*regionStep) % (1.0+regionStep-0.0001);
    }, w, h);
  }
  
  public int regionCount aka nRegions() { ret regionCounter; }
  
  abstract class RegionIterator {
    int pos;
    
    abstract bool next();
    
    int pos() { ret pos; }
    int x() { ret DifferentialRegionsMaker.this.x(pos); }
    int y() { ret DifferentialRegionsMaker.this.y(pos); }
    
    int[] pixelsAsIntArray() { throw todo(); }
  }
  
  // returns points in no particular order
  class FloodRegionIterator > RegionIterator {
    int region;
    new IntBuffer stack; // locations as y*w+x
    BitSet seen = new(size);
    
    *(int *region) {
      int pos = regionFirstPixel.get(region);
      printVars(+region, +pos);
      seen.set(pos);
      stack.add(pos);
    }
    
    // flood-fill region
    bool next() {
      if (empty(stack)) false;
      
      pos = stack.popLast();

      // explore neighborhood
      int x = x(), y = y();
      if (x > 0)   tryPosition(pos-1);
      if (x < w-1) tryPosition(pos+1);
      if (y > 0)   tryPosition(pos-w);
      if (y < h-1) tryPosition(pos+w);

      true;
    }
    
    private void tryPosition(int p) {
      if (!seen.get(p) && regionMatrix[p] == region) {
        seen.set(p);
        stack.add(p);
      }
    }
  }
  
  class CachedRegionIterator > RegionIterator {
    int i, to;

    *(int region) {
      i = regionFirstPixel.get(region);
      to = region+1 < l(regionFirstPixel) ? regionFirstPixel.get(region+1) : l(regionPixels);
      
      ifdef FastRegions_debug
        printVars CachedRegionIterator(+region, +i, +to);
      endifdef
    }
    
    bool next() {
      if (i >= to) false;
      pos = regionPixels.get(i++);
      true;
    }
    
    int[] pixelsAsIntArray() {
      ret regionPixels.subArray(i, to);
    }
  }
  
  int regionSize(int iRegion) { ret regionSize.get(iRegion); }
  
  Pt samplePixel aka firstPixel(int iRegion) {
    ret pt(firstPixelPos(iRegion));
  }
  
  int firstPixelPos(int iRegion) {
    int i = regionFirstPixel.get(iRegion);
    ret regionPixels != null ? regionPixels.get(i) : i;
  }
  
  bool inRegion(int iRegion, int x, int y) {
    ret validPos(x, y) && regionMatrix[pos(x, y)] == iRegion;
  }
  Rect regionBounds(int iRegion) { ret rectFromPoints(
    regionBounds.get((iRegion-1)*4),
    regionBounds.get((iRegion-1)*4+1),
    regionBounds.get((iRegion-1)*4+2),
    regionBounds.get((iRegion-1)*4+3),
  ); }
  
  int regionAt(Pt p) { ret regionAt(p.x, p.y); }
  
  int regionAt(int x, int y) {
    ret !validPos(x, y) ? 0 : regionMatrix[pos(x, y)];
  }
  
  int regionAt(int pos) {
    ret regionMatrix[pos];
  }
  
  RegionIterator regionIterator(int iRegion) {
    ret regionPixels != null
      ? new CachedRegionIterator(iRegion)
      : new FloodRegionIterator(iRegion);
  }
  
  L<Pt> regionPixels(int iRegion) {
    var it = regionIterator(iRegion);
    new L<Pt> l;
    while (it.next())
      l.add(Pt(it.x(), it.y()));
    ret l;
  }
  
  // select extra features before regions are made
  // (not necessary anymore)
  
  void collectFirstPixels { /*regionFirstPixel = new IntBuffer;*/ }
  void collectBounds { /*regionBounds = new IntBuffer;*/ }
  
  BitMatrix regionBitMatrix(int iRegion) {
    ret new AbstractBitMatrix(w, h) {
      public Bool get(int x, int y) {
        ret inRegion(iRegion, x, y);
      }
    };
  }
  
  void markRegionInPixelArray(int[] pixels, int iRegion, int rgba) {
    if (iRegion <= 0) ret;
    for i over pixels:
      if (regionAt(i) == iRegion)
        pixels[i] = rgba;
  }
  
  L<Int> regionIndices() { ret virtualCountList(1, nRegions()+1); }
  
  ItIt<Int> regionsRoughlyByDecreasingSize() {
    ret nestedIterator(countIterator_inclusive_backwards(regionsBySize.size()-1, 0),
      iBucket -> iterator(regionsBySize.get(iBucket)));
  }
  
  public IImageRegion<Img> getRegion aka get(int iRegion) {
    ret new ImageRegion(iRegion);
  }
  
  public L<IImageRegion<Img>> regions aka get() {
    run();
    ret listFromFunction(i -> getRegion(i+1), nRegions());
  }
  
  record noeq ImageRegion(int iRegion) is IImageRegion<Img> {
    public int hashCode() { ret iRegion; }
    public bool equals(O o) {
      if (!o instanceof DifferentialRegionsMaker.ImageRegion) false;
      ret ((DifferentialRegionsMaker.ImageRegion) o).creator() == creator()
        && ((DifferentialRegionsMaker.ImageRegion) o).iRegion == iRegion;
    }
    
    public Img image() { ret image; }
    public O creator() { ret DifferentialRegionsMaker.this; }
    public int indexInCreator() { ret iRegion; }
    public Bool createdWithDiagonals() { ret withDiagonals; }
  
    public Rect bounds() { ret regionBounds(iRegion); }
    
    public int numberOfPixels() { ret regionSize(iRegion); }
    
    public Pt firstPixel() { ret pt(firstPixelPos()); }
    public int firstPixelPos() { ret DifferentialRegionsMaker.this.firstPixelPos(iRegion); }
    
    public ItIt<Pt> pixelIterator() {
      var it = regionIterator(iRegion);
      ret iteratorFromFunction(-> it.next() ? pt(it.pos()) : null);
    }
    
    public int[] pixelsAsIntArray() {
      ret regionIterator(iRegion).pixelsAsIntArray();
    }
    
    public bool contains(int x, int y) { ret inRegion(iRegion, x, y); }
    
    public Color color() {
      ret toColor(rgbForRegion(this));
    }
    
    public int firstPixelRGB() {
      ret getRGB(firstPixelPos());
    }
    
    toString {
      ret renderRecordVars("Region", +brightness(), +color(), pixels := numberOfPixels(), +bounds());
    }
  }
  
  RGB rgbForRegion(ImageRegion r) {
    ret RGB(r.firstPixelRGB());
  }
  
  public Img image() { ret image; }
}