FEnigma/src/fenigma/grid_fix_dialog.py

"""The one modal in the map-import flow: confirm or fix the detected grid.

Nothing else belongs here. Unit detection happens *after* this dialog closes,
because every unit position is expressed in grid coordinates -- detecting
against a grid the user is about to drag would only be thrown away.

The correction handles are the four corners of one cell, not of the whole
screenshot. A homography has 8 degrees of freedom and each dragged corner
contributes 2, so four corners of a single known cell pin it exactly, and a
cell near the frame centre is the one whose corners are easiest to place
accurately by eye. Dragging any handle refits the whole grid immediately, so
the feedback is the entire reconstructed lattice moving, not just a dot.
"""
from __future__ import annotations

import math

import cairo
import gi
import numpy as np

gi.require_version("Gtk", "4.0")
gi.require_version("Adw", "1")

from gi.repository import Adw, Gtk  # noqa: E402

from . import map_vision  # noqa: E402

HANDLE_R = 9.0          # drawn radius of a corner handle, widget px
GRAB_R = 22.0           # how close a press has to be to grab one


def _surface_from_bgr(img):
    """A cairo surface over a numpy BGR image.

    cairo's RGB24 is a 32-bit pixel laid out as B,G,R,x in memory on a
    little-endian machine, which is exactly BGRA, so the converted array can
    back the surface directly with no per-pixel work. The array is kept alive
    by the caller holding it: create_for_data does not copy.
    """
    import cv2

    bgra = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
    bgra = np.ascontiguousarray(bgra)
    h, w = bgra.shape[:2]
    surface = cairo.ImageSurface.create_for_data(
        memoryview(bgra), cairo.FORMAT_RGB24, w, h, w * 4)
    return surface, bgra


class GridFixDialog(Adw.Dialog):
    """Shows the screenshot with the reconstructed grid drawn over it, plus
    four draggable corner handles. on_accept(solution) gets whatever grid is
    on screen when Accept is pressed."""

    def __init__(self, *, image, solution, on_accept, on_discard=None):
        super().__init__(title="Check the detected grid",
                         content_width=900, content_height=760)
        self._image = image
        self._auto = solution
        self._sol = solution
        self._on_accept = on_accept
        self._on_discard = on_discard
        self._surface, self._keepalive = _surface_from_bgr(image)
        self._dragging = None       # index of the handle being dragged
        self._drag_from = None      # its position when the drag began

        quad = map_vision.centre_cell_quad(solution, image.shape)
        # (label, [4 pixel corners], [4 grid corners]); the pixel corners move
        # with the mouse, the grid corners are what they are supposed to BE and
        # never change -- that pairing is the correspondence set refitted from.
        self._label = quad[0] if quad else None
        self._px = list(quad[1]) if quad else []
        self._grid = list(quad[2]) if quad else []

        view = Adw.ToolbarView()
        view.add_top_bar(Adw.HeaderBar())
        self.set_child(view)

        outer = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=10,
                        margin_top=10, margin_bottom=10, margin_start=10, margin_end=10)

        self._area = Gtk.DrawingArea(vexpand=True, hexpand=True)
        self._area.set_draw_func(self._draw)
        drag = Gtk.GestureDrag()
        drag.connect("drag-begin", self._on_drag_begin)
        drag.connect("drag-update", self._on_drag_update)
        drag.connect("drag-end", lambda *_a: setattr(self, "_dragging", None))
        self._area.add_controller(drag)
        outer.append(self._area)

        cell = self._label or "?"
        self._hint = Gtk.Label(xalign=0, css_classes=["dim-label"], wrap=True)
        self._hint.set_label(
            f"Grid solved from {solution.votes} label read(s). "
            f"If it is off, drag the four handles onto the corners of cell {cell}."
            if self._px else
            f"Grid solved from {solution.votes} label read(s)."
        )
        outer.append(self._hint)

        buttons = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=8,
                          halign=Gtk.Align.END)
        discard = Gtk.Button(label="Discard", css_classes=["pill"])
        discard.connect("clicked", lambda _b: self._discard())
        buttons.append(discard)
        if self._px:
            reset = Gtk.Button(label="Reset", css_classes=["pill"],
                               tooltip_text="Back to the automatically detected grid")
            reset.connect("clicked", lambda _b: self._reset())
            buttons.append(reset)
        accept = Gtk.Button(label="Use this grid",
                            css_classes=["pill", "suggested-action"])
        accept.connect("clicked", lambda _b: self._accept())
        buttons.append(accept)
        outer.append(buttons)

        view.set_content(outer)

    # -- geometry ------------------------------------------------------------
    def _fit(self):
        """(scale, ox, oy) letterboxing the screenshot into the drawing area."""
        w, h = self._area.get_width(), self._area.get_height()
        ih, iw = self._image.shape[:2]
        if not w or not h:
            return 1.0, 0.0, 0.0
        s = min(w / iw, h / ih)
        return s, (w - iw * s) / 2, (h - ih * s) / 2

    def _to_widget(self, p):
        s, ox, oy = self._fit()
        return p[0] * s + ox, p[1] * s + oy

    def _to_image(self, x, y):
        s, ox, oy = self._fit()
        return (x - ox) / s, (y - oy) / s

    def _refit(self):
        """Rebuild the grid from the four handle positions.

        A bad drag (two handles on top of each other) makes the homography
        degenerate; keep the previous grid rather than crash, the next drag
        update recovers.
        """
        try:
            self._sol = map_vision.solution_from_correspondences(
                list(zip(self._grid, self._px)))
        except (ValueError, np.linalg.LinAlgError):
            pass
        self._area.queue_draw()

    def _reset(self):
        quad = map_vision.centre_cell_quad(self._auto, self._image.shape)
        if quad:
            self._px = list(quad[1])
        self._sol = self._auto
        self._area.queue_draw()

    # -- input ---------------------------------------------------------------
    def _on_drag_begin(self, _gesture, x, y):
        self._dragging = None
        best = GRAB_R
        for i, p in enumerate(self._px):
            wx, wy = self._to_widget(p)
            d = ((wx - x) ** 2 + (wy - y) ** 2) ** 0.5
            if d < best:
                best, self._dragging = d, i
        if self._dragging is not None:
            self._drag_from = self._px[self._dragging]

    def _on_drag_update(self, _gesture, dx, dy):
        if self._dragging is None:
            return
        s, _ox, _oy = self._fit()
        if s <= 0:
            return
        fx, fy = self._drag_from
        self._px[self._dragging] = (fx + dx / s, fy + dy / s)
        self._refit()

    def _accept(self):
        self.close()
        self._on_accept(self._sol)

    def _discard(self):
        self.close()
        if self._on_discard is not None:
            self._on_discard()

    # -- drawing -------------------------------------------------------------
    def _draw(self, _area, cr, width, height):
        cr.set_source_rgb(0.08, 0.08, 0.08)
        cr.paint()
        s, ox, oy = self._fit()

        cr.save()
        cr.translate(ox, oy)
        cr.scale(s, s)
        cr.set_source_surface(self._surface, 0, 0)
        cr.get_source().set_filter(cairo.FILTER_GOOD)
        cr.paint()
        cr.restore()

        self._draw_grid(cr)

        for i, p in enumerate(self._px):
            wx, wy = self._to_widget(p)
            # new_path() before every arc: cairo's arc() joins the current point
            # to the arc's start, and _draw_grid leaves one behind at the last
            # cell name it drew. Without this, the first handle gets a stray
            # line reaching across the whole screenshot from that label.
            cr.new_path()
            cr.set_source_rgb(1.0, 0.85, 0.1)
            cr.arc(wx, wy, HANDLE_R, 0, 2 * math.pi)
            cr.fill_preserve()
            cr.set_source_rgb(0.1, 0.1, 0.1)
            cr.set_line_width(2.0)
            cr.stroke()
            if i == self._dragging:
                cr.new_path()
                cr.set_source_rgb(1.0, 1.0, 1.0)
                cr.arc(wx, wy, HANDLE_R + 4, 0, 2 * math.pi)
                cr.set_line_width(1.5)
                cr.stroke()

    def _draw_grid(self, cr):
        """Every in-range cell the grid puts inside the frame, with its name
        drawn where the game draws it. A wrong grid is obvious precisely
        because those names land off the painted labels."""
        sol = self._sol
        h, w = self._image.shape[:2]
        inv = np.linalg.inv(sol.H)
        corners = inv @ np.array([[0, w, w, 0], [0, 0, h, h], [1, 1, 1, 1.0]])
        if np.any(np.abs(corners[2]) < 1e-9):
            return
        ij = corners[:2] / corners[2]
        cr.set_line_width(1.6)
        cr.select_font_face("Sans", cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
        L2G = sol.lattice_to_grid()
        for i in range(int(np.floor(ij[0].min())) - 1, int(np.ceil(ij[0].max())) + 2):
            for j in range(int(np.floor(ij[1].min())) - 1, int(np.ceil(ij[1].max())) + 2):
                g = L2G @ np.array([i, j, 1.0])
                col, row = int(round(g[0])), int(round(g[1]))
                if not (0 <= col < map_vision.COLS and 1 <= row <= map_vision.ROWS):
                    continue
                quad = sol.H @ np.array([[i, i + 1, i + 1, i],
                                         [j, j, j + 1, j + 1], [1, 1, 1, 1.0]])
                if np.any(np.abs(quad[2]) < 1e-9):
                    continue
                pts = [self._to_widget(p) for p in (quad[:2] / quad[2]).T]
                cr.new_path()
                cr.set_source_rgba(1.0, 1.0, 0.2, 0.75)
                cr.move_to(*pts[0])
                for p in pts[1:]:
                    cr.line_to(*p)
                cr.close_path()
                cr.stroke()

                # The game pads a cell's label in from its top-left corner by a
                # fixed fraction of the cell, which is also how the solver finds
                # labels in the first place (see map_vision.PAD_L/PAD_T).
                lx = i + (map_vision.PAD_L if sol.si > 0 else 1 - map_vision.PAD_L)
                ly = j + (map_vision.PAD_T if sol.sj > 0 else 1 - map_vision.PAD_T)
                t = sol.H @ np.array([lx, ly, 1.0])
                if abs(t[2]) < 1e-9:
                    continue
                tx, ty = self._to_widget((t[0] / t[2], t[1] / t[2]))
                side = float(np.hypot(pts[1][0] - pts[0][0], pts[1][1] - pts[0][1]))
                cr.set_font_size(max(9.0, min(30.0, side * 0.16)))
                name = f"{map_vision.LARGE_X[col]}{row}"
                cr.move_to(tx, ty)
                cr.set_source_rgba(0, 0, 0, 0.8)
                cr.text_path(name)
                cr.set_line_width(3.0)
                cr.stroke()
                cr.move_to(tx, ty)
                cr.set_source_rgb(0.3, 1.0, 0.3)
                cr.show_text(name)