!7

///////////////////////////
// Your API to work with //
///////////////////////////

abstract sclass GameForAI {
  abstract RGBImage getImage();
  abstract L<Rect> submit(Pt point); // returns correct solution so AI can learn from it
}

abstract sclass AI {
  // stuff you get
  
  GameForAI game;
  RGBImage image;
  int steps;
  RGBImage image() { ret image; }
  int w() { ret image.w(); }
  int h() { ret image.h(); }
  L<Rect> submit(Pt p) { ret game.submit(p); }
  L<Rect> submit(Rect r) { ret submit(middleOfRect(r)); }
  L<Rect> submitNothing() { ret game.submit(null); }

  // implement this method and call submit()
  abstract void go();
  
  void done() {}
}

//////////////////////////////////////
// Test AIs. Just add your own here //
//////////////////////////////////////

AI > ClickAnyNumberAI {
  void go {
    submit(random(segment(image;
  }
}

AI > Homos {
  static ImageSurface is;
  void go {
    L<Rect> solution = submitNothing();
    if (l(solution) == 2)
      is = showImage(is, "Homos!", mergeImagesHorizontally(rgbClips(image, solution)));
  }
}

AI > Heteros {
  static ImageSurface is;
  void go {
    L<Rect> solution = submitNothing();
    if (l(solution) == 1) {
      L<Rect> myRects = segment(image);
      Rect sol = autoCropOfRGBImage(image, first(solution));
      myRects = dropRectsIntersectingWith(myRects, sol);
      if (l(myRects) == 1)
        is = showImage(is, "Heteros!", mergeImagesHorizontally(rgbClips(image, concatLists(ll(sol), myRects)));
    }
  }
}

AI > AnalyzeFeatureExtractor {
  O extractor; // func(RGBImage) -> anything comparable with eq()
  int ok;
  
  *() {
    extractor = func(RGBImage img) { rgbMD5(img) };
  }
  
  void go {
    L<Rect> solution = submitNothing();
    Analysis analysis = analyzePair(image, solution);
    if (analysis == null) ret;
    bool same = eq(
      pcallF(extractor, first(analysis.images)),
      pcallF(extractor, second(analysis.images)));
    if (same == !analysis.hetero) ++ok;
  }
  
  void done() {
    print("OK: " + ok + " of " + steps);
  }
}

AI > AnalyzeCategorizer {
  new ImageCategorizer cat;
  int ok;
  static ImageSurface is;
  
  RGBImage simplifyImage(RGBImage img) {
    ret new Image2B(img).toRGB();
  }
  
  Rect guess() {
    L<Rect> myRects = segment(image);
    if (l(myRects) != 2) null;
    int i = cat.locateImage(simplifyImage(image.clip(first(myRects;
    int j = cat.locateImage(simplifyImage(image.clip(second(myRects;
    ret get(myRects, i < j ? 0 : 1);
  }
  
  void go {
    L<Rect> solution = submit(guess());
    Analysis analysis = analyzePair(image, solution);
    if (analysis == null) ret;
    int n = cat.numCategories();
    int i = cat.addImage(simplifyImage(first(analysis.images)));
    int j = cat.addImage(simplifyImage(second(analysis.images)));
    if ((i == j) == !analysis.hetero) ++ok;
    bool update = cat.numCategories() > n;
    if (analysis.hetero && i > j) {
      swap(cat.images, i, j);
      update = true;
    }
    if (cat.numCategories() > 10) {
      cat.allowedDistance += 0.01f;
      print("Got 11 categories after " + steps + " rounds. Upping allowed distance to: " + cat.allowedDistance);
      cat.clear();
      steps = ok = 0;
    } else if (update) {
      bool first = is == null;
      is = showImage(is, "Categories!", mergeImagesHorizontally(cat.images));
      if (first) {
        setFrameWidth(is, 400);
        moveToTopRightCorner(getFrame(is));
      }
    }
  }
  
  void done() {
    print("OK: " + ok + " of " + steps + ", categories: " + cat.numCategories() + " (should probably be 10)");
  }
}

sclass Analysis {
  bool hetero;
  L<RGBImage> images;
  
  *() {}
  *(bool *hetero, L<RGBImage> *images) {}
}

static Analysis analyzePair(RGBImage image, L<Rect> solution) {
  if (l(solution) == 2)
    ret new Analysis(false, rgbClips(image, rgbAutoCropRects(image, solution)));
  if (l(solution) == 1) {
    L<Rect> myRects = segment(image);
    Rect sol = autoCropOfRGBImage(image, first(solution));
    myRects = dropRectsIntersectingWith(myRects, sol);
    if (l(myRects) == 1)
      ret new Analysis(true, rgbClips(image, concatLists(ll(sol), myRects)));
  }
  null; // no insights today
}

static L<Rect> segment(RGBImage img) {
  ret autoSegment(new BWImage(img), 3); // 2 is too small
}

///////////////
// Main Game //
///////////////

static int w = 150, h = 150;
static int border = 10, fontSize = 30, numberMargin = 2;
static int rounds = 1000, numbers = 2;
static S font = #1004887;
static bool alwaysNewImage = true; // prevents AIs being stuck on an image
static int fps = 50;

static int points, clicks;
static ImageSurface is, isWrong;
static RGBImage img;
static new HashMap<Rect, Int> words;
static L<Rect> solution;
volatile sbool aiMode;
static JLabel lblScore, lblTiming;
static new HashMap<Class, AI> ais;

p-substance {
  nextImage();
  addToWindow(is, withMargin(lblScore = jcenteredBoldLabel("Your score: 0 of 0")));
  for (final Class c : myNonAbstractClassesImplementing(AI)) {
    final S name = shortClassName(c);
    final JButton btn = jbutton(name);
    onClick(btn, r {
      if (aiMode) ret;
      btn.setText(name + " Running...");
      thread {
        Game game = testAI(c, voidfunc(Game game) {
          if ((game.step % 10) == 0 || game.step == rounds) {
            S text = name + " scored " + game.points + " of " + game.step;
            if (game.error != null)
              text += " - ERROR: " + game.error;
            setText_noWait(btn, text);
          }
        });
      }
    });
    addToWindow(is, withMargin(btn));
  }
  addToWindow(is, lblTiming = jCenteredLabel(;
  titlePopupMenuItem(is, "Restart AIs", r { ais.clear() });
  addToWindow(is, withMargin(jbutton("Restart AIs", r { ais.clear(); })));
  
  packFrame(is);
  setFrameWidth(is, 400);
  centerTopFrame(is);
  //hideConsole();
}

svoid nextImage {
  RGBImage img = rgbImage(Color.white, w, h);
  words.clear();
  new L<Rect> taken;
  for i to numbers: {
    int num = random(10);
    S s = str(num);
    RGB color = new RGB(random(0.5), random(0.5), random(0.5));
    renderText_fg.set(color.getColor());
    BufferedImage ti = renderText(font, fontSize, s);
    Rect r;
    int safety = 0;
    do {
      r = randomRect(w, h, border, ti.getWidth(), ti.getHeight());
      if (r == null || ++safety >= 1000) fail("Image too small: \*ti.getWidth()*/*\*ti.getHeight()*/ > \*w*/*\*h*/");
    } while (anyRectOverlaps(taken, r));
    rgbCopy(ti, img, r.x, r.y);
    words.put(r, num);
    taken.add(growRect(r, numberMargin));
  }
  solution = keysWithBiggestValue(words);
  
  main.img = img;
  if (aiMode) ret; // occasional updates only when AI is running
  showTheImage();
}

svoid showTheImage {
  bool first = is == null;
  is = showImage_centered(is, img, "Click On The Highest Number!");
  if (first) {
    onLeftClick(is, voidfunc(Pt p) {
      ++clicks;
      if (anyRectContains(solution, p)) {
        ++points;
        nextImage();
      } else if (alwaysNewImage) nextImage();
      lblScore.setText(print("Your score: " + points + " of " + clicks));
    });
    disableImageSurfaceSelector(is);
    
    // animate if idle
    awtEvery(is, 1000, r {
      if (!aiMode && isInForeground(is) && !mouseInComponent(is))
        nextImage();
    });
    
    // show current image occasionally when AI is running
    awtEvery(is, 1000/fps, r {
      if (aiMode) showTheImage();
    });
  }
}

// AI stuff

sclass Game extends GameForAI {
  int points, step;
  bool submitted;
  Throwable error;
  
  RGBImage getImage() { ret img; }

  L<Rect> submit(Pt p) {
    if (submitted) fail("No multi-submit please");
    submitted = true;
    if (p != null && anyRectContains(solution, p)) {
      ++points;
      if (!alwaysNewImage) nextImage();
    }
    ret solution;
  }
}

static Game scoreAI(AI ai, long rounds, O onScore) {
  aiMode = true;
  final long start = sysNow();
  try {
    new Game game;
    setOpt(ai, +game);
    while (game.step < rounds) {
      ++game.step;
      ++ai.steps;
      game.submitted = false;
      int points = game.points;
      RGBImage image = img;
      try {
        setOpt(ai, +image);
        ai.go();
      } catch e {
        if (game.error == null) { // print first error to console
          showConsole();
          printStackTrace(e);
        }
        game.error = e;
      } finally {
        setOpt(ai, image := null);
      }
      if (points == game.points) {
        bool first = isWrong == null;
        isWrong = showImage(isWrong, "Blunder!", image);
        if (first) {
          setFrameWidth(isWrong, 200);
          moveToTopRightCorner(isWrong);
          moveFrameDown(isWrong, 100);
        }
      }
      pcallF(onScore, game);
      //print("Step " + step + ", move: " + p + ", points: " + game.points);
      if (alwaysNewImage) nextImage();
    }
    print("AI " + shortClassName(ai) + " points after " + rounds + " rounds: " + game.points);
    ai.done();
    ret game;
  } finally {
    aiMode = false;
    awt {
      lblTiming.setText(formatDouble(toSeconds(sysNow()-start), 1) + " s");
      nextImage();
    }
  }
}

static AI getAI(Class<? extends AI> c) {
  AI ai = ais.get(c);
  if (ai == null) ais.put(c, ai = nu(c));
  ret ai;
}

// onScore: voidfunc(Game)
static Game testAI(Class<? extends AI> c, O onScore) {
  ret scoreAI(getAI(c), rounds, onScore);
}