!7

sbool real = true, test = false;
static double threshold = 0.5;

sS data = [[
  RecognizedText(r=Rect(h=12, w=54, x=627, y=309), text=Scored(score=fl "1.0", v="Confirm"))
  RecognizedText(r=Rect(h=12, w=36, x=684, y=309), text=Scored(score=fl "1.0", v="close"))
  RecognizedText(r=Rect(h=10, w=23, x=494, y=359), text=Scored(score=fl "0.9542266", v="you"))
  RecognizedText(r=Rect(h=11, w=39, x=548, y=358), text=Scored(score=fl "1.0", v="about"))
  RecognizedText(r=Rect(h=10, w=13, x=592, y=359), text=Scored(score=fl "0.9918628", v="to"))
  RecognizedText(r=Rect(h=11, w=33, x=610, y=358), text=Scored(score=fl "0.998467", v="close"))
  RecognizedText(r=Rect(h=10, w=6, x=648, y=359), text=Scored(score=fl "0.8560784", v="3"))
  RecognizedText(r=Rect(h=11, w=31, x=659, y=358), text=Scored(score=fl "1.0", v="tabs."))
  RecognizedText(r=Rect(h=13, w=50, x=695, y=359), text=Scored(score=fl "0.9463429", v="Areyou"))
  RecognizedText(r=Rect(h=10, w=33, x=812, y=359), text=Scored(score=fl "0.9773897", v="want"))
  RecognizedText(r=Rect(h=10, w=13, x=850, y=359), text=Scored(score=fl "0.9918628", v="to"))
  RecognizedText(r=Rect(h=11, w=65, x=868, y=358), text=Scored(score=fl "1.0", v="continue?"))
  RecognizedText(r=Rect(h=8, w=20, x=523, y=361), text=Scored(score=fl "0.97618467", v="are"))
  RecognizedText(r=Rect(h=8, w=28, x=751, y=361), text=Scored(score=fl "0.9821385", v="sure"))
  RecognizedText(r=Rect(h=11, w=23, x=783, y=361), text=Scored(score=fl "0.9744504", v="you"))
  RecognizedText(r=Rect(h=11, w=43, x=765, y=406), text=Scored(score=fl "0.99504626", v="Cancel"), clickTarget="Cancel")
  RecognizedText(r=Rect(h=11, w=35, x=855, y=406), text=Scored(score=fl "1.0", v="Close"), clickTarget="Close tabs")
  RecognizedText(r=Rect(h=11, w=28, x=895, y=406), text=Scored(score=fl "0.9762714", v="tabs"), clickTarget="Close tabs")
]];

p {
  L<RecognizedText> expected = map unstruct(tlft(data));
  
  if (test)
    selfTest(expected);
  
  if (real) {
    print("\nLive mode!\n");
    mouseMover().endAfterInterference = false;
    repeat with sleep 1 {
      if (mouseMover().active2()) {
        Pair<Double, Pt> p = recognize(expected, shootScreen2());
        if (p != null && p.a >= threshold) {
          print("Found dialog at " + p.b);
          L<Pt> targets = getClickTargets(expected, p.b).get("Cancel");
          mouseMover().moveMouse(oneOf(targets));
        } else print("No dialog found");
      }
    }
  }
}

svoid selfTest(L<RecognizedText> expected) {  
  bool failures = false;
  for (S imageID, Bool containsDialog : litorderedmap(#1101350, true, #1101351, true, #1101320, false)) {
    Pair<Double, Pt> p = recognize(expected, loadImage2(imageID));
    print(containsDialog + " / " + p);
    
    bool found = p != null && p.a >= threshold;
    if (containsDialog != found) { print("detection failure"); failures = true; }
    else if (found) {
      print("Found dialog correctly");
      //MultiMap<S, Pt> targets = getClickTargets(expected, p.b);
    }
  }
}

static MultiMap<S, Pt> getClickTargets(L<RecognizedText> expected, Pt shift) {
  new MultiMap<S, Pt> mm;
  for (RecognizedText t : expected) {
    S ct = getString(t, 'clickTarget);
    if (ct != null) {
      Pt pt = ptAdd(rectCenter(t.r), shift);
      print("Click target: " + ct + " => " + pt);
      mm.put(ct, pt);
    }
  }
  ret mm;
}

// returns (closeness, shift)
static Pair<Double, Pt> recognize(L<RecognizedText> expected, BufferedImage img) {
  L<RecognizedText> actual = ocr_recognizeMultiLine_scored(img);
  
  MultiMap<S, RecognizedText> byWord = multiMapIndexOverMethod(actual, 'text);
  
  new Best<Pt> bestShift;
  bestShift.put(pt(0, 0), tryShift(expected, byWord, pt(0, 0)));
  
  Map<Pt, Int> shifts = multiSetAsMap_popularFirst(possibleShifts(expected, byWord));
  //pnlStruct(shifts);
  Pt shift = firstKey(shifts);
  bestShift.put(shift, tryShift(expected, byWord, shift));
  ret best_reversedPair(bestShift);
}

static MultiSet<Pt> possibleShifts(L<RecognizedText> expected, MultiMap<S, RecognizedText> byWord) {
  new MultiSet<Pt> ms;
  for (RecognizedText t : expected) {
    L<RecognizedText> list = byWord.get(t.text());
    if (empty(list)) continue;
    final Pt p = rectCenter(t.r);
    for (RecognizedText r : list)
      ms.add(ptDiff(rectCenter(r.r), p));
  }
  ret ms;
}

static double tryShift(L<RecognizedText> expected, MultiMap<S, RecognizedText> byWord, Pt shift) {
  double closeness = 0;
  for (RecognizedText t : expected) {
    L<RecognizedText> list = byWord.get(t.text());
    //print(t.text() + " " + t.r + " => " + struct(collect('r, list)));
    if (empty(list)) continue;
    final Pt p = ptAdd(rectCenter(t.r), shift);
    double dist = sqrt(minOfDoubles(map(list, func(RecognizedText r) {
      ptDistanceSquared(p, rectCenter(r.r))
    })));
    closeness += 1/(dist+1);
    //print(t.text() + " | dist=" + dist + ", c=" + closeness);
  }
  closeness /= l(expected);
  //print("Closeness: " + closeness + " (" + l(expected) + ")");
  ret closeness;
}