!7

cmodule ChessBoardRecognizer {

  S segmenterPreset;

  ParameterizedSegmenter segmenter;

  L<Rect> segments;

  Rect boardLocation;

  double recognitionScore;

  S status;

  S selectedSquare;

  S piece;

  LPair<S, Double> recognized;

  S fen;

  switchable bool squareHeightAtBottom = true; // for YouTube videos

  transient int recreatedSquareSize = 32;

  transient ReliableSingleThread rst = dm_rst(module(), r recognize);

  transient BufferedImage img, board, square, recreated;

  transient LL<BufferedImage> squareImages;

  transient ImageSurface isInput, isBoard, isSquare, isRecreated;

  transient LS pieces = ai_chessPieces();

  transient Map<Bool, Map<S, BufferedImage>> recreateStock; // false = dark, true = light

  visualize {

    isBoard = jImageSurface();

    imageSurfaceOnHover(isBoard, voidfunc(Pt p) {

      if (board == null || p == null) ret with noToolTip(isBoard);

      setToolTip(isBoard, toolTipForSquare(boardPointToSquare(p));

    });

    imageSurfaceOnLeftClick(isBoard, voidfunc(Pt p) {

      selectSquare(boardPointToSquare(p))

    });

    isRecreated = jImageSurface(recreated);

    imageSurfaceOnHover(isRecreated, voidfunc(Pt p) {

      if (recreated == null || p == null) ret with noToolTip(isRecreated);

      setToolTip(isRecreated, toolTipForSquare(recreatedPointToSquare(p));

    });

    imageSurfaceOnLeftClick(isRecreated, voidfunc(Pt p) {

      selectSquare(recreatedPointToSquare(p))

    });

    ret withCalc(rst, jvsplit(

      jCenteredSection("Input image", jscroll_center(isInput = jImageSurface(or_func(img, () -> whiteImage(100))))),

      northCenterAndSouthWithMargins(

        centerAndEastWithMargin(withLabel("Segmenter preset: ", main.onChange(rst, dm_comboBox segmenterPreset(agiBlue_chessBoardSegmenterPresetNames()))),

          hstackWithSpacing(

            jbutton("Recognize", rst),

            jPopDownButton_noText(

              "Show possible chess boards found", rThread showSegments,

              "Publish training image...", rThread publishImage,

              "---", null,

              "Screenshot (hiding OS)", rThread {

                isInput.setImageAndZoomToDisplay(dm_shootScreenHidingOS());

                rst.trigger();

              })

          )),

        hgridWithSpacing(

          jCenteredSection("Recreated", jscroll_center(isRecreated)),

          jCenteredSection("Chess board in image",

            centerAndSouthWithMargin(jscroll_center(isBoard),

            dm_centeredCalculatedLabel(() -> "Recognition score: " + formatDouble(recognitionScore, 2))

          )),

          jCenteredLiveValueSection(dm_calculatedLiveValue(S, () -> "Selected square" + (empty(selectedSquare) ? "" : " (" + selectedSquare + ")")), centerAndSouthWithMargin(

            jscroll_center(isSquare = jImageSurface(square)),

            rightAlignedLine(

              withLabel("Piece: ", dm_comboBox('piece, pieces)),

              jbutton("Save & upload", rThread saveAndUploadPiece))

        ))),

        vstackWithSpacing(

          dm_label('status),

          withLabel("FEN:", dm_label('fen))

        )

      )));

  }

  void recognize {

    status("Recognizing...");

    segmenter = parameterizedSegmenterFromAGIBlue(segmenterPreset);

    if (segmenter == null) ret with infoBox(status("No segmenter"));

    //img = dm_shootScreenHidingOS();

    if (isInput != null) img = isInput.getImage();

    L<Rect> rects = segmenter.get(img);

    setField(segments := rects);

    if (empty(rects)) {

      clearImageSurfaces(isBoard, isRecreated);

      setField(fen := "");

      ret with status("No chess board found");

    }

    Rect r = first(rects);

    bool corrected = squareHeightAtBottom && abs(r.w-r.h) > 3;

    if (corrected) r.h = r.w;

    new LS comments;

    if (l(rects) > 1) comments.add("randomly choosing one from " + nBoards(rects));

    if (corrected) comments.add("height-corrected");

    status("Chess board found at: " + r + appendRoundBracketed(joinWithComma(comments)));

    setField(boardLocation := r);

    board = clipBufferedImage(img, r);

    squareImages = bufferedImageMNGrid(board, 8, 8);

    if (isBoard != null)

      isBoard.setImageAndZoomToDisplay(

        mergeBufferedImagesVertically(

          map(bufferedImageNVerticalSlices(8, board), i ->

            mergeBufferedImagesHorizontally(bufferedImageNHorizontalSlices(8, i)))));

    new ChessPieceRecognizer pieceRecognizer;

    pieceRecognizer.load();

    recognized = new L;

    for (BufferedImage img : concatLists(squareImages))

      recognized.add(pieceRecognizer.recognize(img));

    setField(recognitionScore := doubleAvg(pairsB(recognized));

    setField(fen := makeFEN());

    print(recognized);

    recreate();

  }

  S status(S status) { setField(+status); ret status; }

  O _getReloadData() { ret img; }

  void _setReloadData(BufferedImage img) { this.img = img; }

  void setImage(BufferedImage img) {

    this.img = img;

    if (isInput != null) isInput.setImageAndZoomToDisplay(img);

    rst.trigger();

  }

  S boardPointToSquare(Pt p) {

    double space = imageMergeSpacing();

    double squareW = board.getWidth()/8.0+space;

    double squareH = board.getHeight()/8.0+space;

    int x = clamp(iratio_floor(p.x+space/2, squareW), 0, 7);

    int y = clamp(iratio_floor(p.y+space/2, squareH), 0, 7);

    ret strCharPlus('A', x) + (8-y);

  }

  S recreatedPointToSquare(Pt p) {

    int x = clamp(iratio_floor(p.x, recreatedSquareSize), 0, 7);

    int y = clamp(iratio_floor(p.y, recreatedSquareSize), 0, 7);

    ret strCharPlus('A', x) + (8-y);

  }

  void saveAndUploadPiece enter {

    if (square == null || empty(piece)) ret;

    chessOCR_uploadPieceImage(selectedSquare, piece, square);

    if (mapGet(mapGet(recreateStock, chess_isLightSquare(selectedSquare)), piece) == null) recreateStock = null;

    rst.trigger();

  }

  void recreate {

    if (recognized == null) ret;

    if (recreateStock == null)

      recreateStock = mapToValues(ll(false, true), light ->

        mapValues(pairsToMap(reversePairs(chessPieceImagesFromAGIBlue(light))), imageID -> resizeImage(loadImage2(imageID), recreatedSquareSize, recreatedSquareSize)));

    new L<BufferedImage> out;

    int i = 0;

    for (S piece : pairsA(recognized)) {

      bool light = even(i/8+i%8);

      out.add(or_func(mapGet(mapGet(recreateStock, light), piece), () -> whiteImage(recreatedSquareSize, recreatedSquareSize)));

      ++i;

    }

    new L<BufferedImage> rows;

    for (L<BufferedImage> row : listToChunks(out, 8))

      rows.add(mergeBufferedImagesHorizontally(row, spacing := 0));

    recreated = mergeBufferedImagesVertically(rows, spacing := 0);

    if (isRecreated != null) isRecreated.setImageAndZoomToDisplay(recreated);

  }

  Pair<S, Double> recognitionForSquare(S square) {

    int x, y = unpair chess_squareToPos(square);

    ret _get(recognized, y*8+x);

  }

  S toolTipForSquare(S square) {

    Pair<S, Double> rec = recognitionForSquare(square);

    ret square + " - " + (rec == null ? "?" : rec.a + " [" + iround(clamp(rec.b, 0, 100)) + "%]");

  }

  void selectSquare(S _square) {

    setField(selectedSquare := _square);

    S rec = pairA(recognitionForSquare(selectedSquare));

    if (rec != null) setField(piece := rec);

    int x, y = unpair chess_squareToPos(selectedSquare);

    square = _get(_get(squareImages, y), x);

    ChessPieceProfile1 profile = chessOCR_pieceProfileFromRawImage_1(square, withPreprocessedImages := true);

    isSquare.setImage(mergeBufferedImagesHorizontally(

      itemPlusList(square, map toBufferedImage(profile.preprocessedImages))));

    //printVars_str("leftClick", +p, +selectedSquare, +x, +y, +square);

  }

  void showSegments {

    L<Rect> segments = this.segments;

    showImage_centered(n2(segments, "segment") + " found",

      mergeBufferedImagesVertically(map(segments, r -> clipBufferedImage(img, r))));

  }

  void publishImage enter {

    BufferedImage img = this.img;

    S segmenterPreset = this.segmenterPreset;

    JCheckBox cbFound = jcheckbox(true), cbHasBoard = jcheckbox(true);

    JTextField tfFEN = jtextfield(fen), tfComments = jtextfield(),

      tfImageCredits = jtextfield();

    showFormTitled("Publish training image",

      "Image", jImageSurface(scaleImageToWidth(img, 500)),

      "Segmenter used", jlabel(segmenterPreset),

      "Is there a chess board in the image?", cbHasBoard,

      "FEN (empty if no board in image)", tfFEN,

      "Was the chess board found?", cbFound, // TODO: make clearer what this means if there is no board in image

      "Image credits (optional)", tfImageCredits,

      "Random comments", tfComments,

      "NOTE:", jMultiLineLabel("By clicking 'publish', you PUBLISH this image (fair use/public domain)."),

      "", jThreadedButton("Publish", r {

        temp tempInfoBox_noHide("Uploading training image...");

        S imageURL = uploadToImageServer("Chess board recognition training image", img);

        S desc = "Chess board recognition training image " + assertNempty("Parseable image URL", parseSnippetImageURL(imageURL));

        new LinkedHashMap<S> map;

        map.put("Input image", imageURL);

        mapPutIfNemptyValue(map, "Image credits", gtt(tfImageCredits));

        S fen = gtt(tfFEN);

        map.put("Image has chess board", yesNo_short_firstUpper(isChecked(cbHasBoard) || nempty(fen)));

        mapPutIfNemptyValue(map, "FEN", fen);

        map.put("Segmenter used", segmenterPreset);

        map.put("Chess board found", yesNo_short_firstUpper(isChecked(cbFound)));

        if (isChecked(cbFound))

          map.put("Board location", struct(boardLocation));

        S sliceID = agiBlue_chessBoardRecognitionTrainingImagesSliceID();

        new L<Map> toPost;

        S name = agiBlue_createUnusedNumberedPage(sliceID, desc + " #");

        for (S key, value : map)

          toPost.add(litmap(q := name, +key, +value));

        agiBot_postMulti(keyPairForProgram(), toPost, slice := sliceID);

        infoBox("Training image uploaded!");

        openURLInBrowser(agiBlue_linkForPhrase(name, slice := sliceID));

        disposeWindow(heldButton());

      })

    );

  }

  S makeFEN() {

    ret chess_makeFEN(listToChunks(pairsA(recognized), 8));

  }

}