calibrate.py

"""Interactive calibration for every bot in this repo."""

import json
import os
import re
import shutil
import time

import keyboard
import pyautogui

import config

CAPTURE_KEY = "f8"
SKIP_KEY = "f9"
QUIT_KEY = "esc"

# click targets, in the order the calibrator walks through them
POINTS = [
    ("respawn_button", "the respawn button in the top-right of the HUD"),
    ("respawn_confirm", "the confirm button in the respawn dialog"),
    ("tool_slot", "the inventory slot the bots toggle between wrench and fist"),
    ("inventory_drag_target", "where the inventory panel gets dragged to when opened"),
    ("inventory_scroll_point", "a point inside the open inventory list, for scrolling"),
    ("wrench_marker_object", "the object wrench-marker.py wrenches on each level"),
    (
        "wrench_marker_last_object",
        "the object wrench-marker.py wrenches on the final level",
    ),
    ("wrench_marker_ok", "the OK button in the marker dialog"),
    ("wrench_portal_object", "the portal wrench-portal.py wrenches on each level"),
    ("wrench_portal_text_field", "the destination text field in the portal dialog"),
    ("wrench_portal_ok", "the OK button in the portal dialog"),
]

# pixel probes. 'color' records the RGB, 'brightness' only the position
PROBES = [
    ("portal", "color", "a pixel inside the level portal, used to count levels"),
]

# bands the character must stand within to reach a vend
BANDS = [
    ("block_vend", "the block vending machine"),
    ("seed_vend", "the seed vending machine"),
]

# templates for locateCenterOnScreen, re-crop when the UI scale changes
ASSETS = [
    (
        "gt-body.png",
        (
            "your character's body, crop a part that looks the same "
            "facing either way, it is matched regardless of direction"
        ),
    ),
    ("inventory.png", "the inventory button / panel handle"),
    ("wrench-inventory.png", "the wrench sitting in the inventory slot"),
    ("chand.png", "a chandelier block in the inventory"),
    ("chand-200.png", "a chandelier block stack showing 200"),
    ("chand-seed.png", "a chandelier seed in the inventory"),
    ("chand-seed-200.png", "a chandelier seed stack showing 200"),
    ("ftank.png", "a fish tank block in the inventory"),
    ("ftank-200.png", "a fish tank block stack showing 200"),
    ("ftank-seed.png", "a fish tank seed in the inventory"),
    ("ftank-seed-200.png", "a fish tank seed stack showing 200"),
    ("put-item.png", 'the "Put item to machine" button'),
    ("add-machine.png", 'the "Add to machine" button'),
    ("empty-machine.png", 'the "Empty the machine" button'),
    ("update-vend.png", 'the "Update" button in the vend dialog'),
]

BACKUP_DIR = os.path.join(config.ASSETS_DIR, "_backup")


class Quit(Exception):
    pass


def wait_for_key(keys):
    """Block until one of `keys` is pressed, then released."""
    while True:
        for key in keys:
            if keyboard.is_pressed(key):
                while keyboard.is_pressed(key):
                    time.sleep(0.02)
                return key
        time.sleep(0.02)


def capture_point(label, description, current):
    """Ask for one screen position, or None to keep `current`."""
    # format up front, a tuple would expand into args
    shown = "not set" if not current else f"({current[0]}, {current[1]})"

    print(f"\n  {label}")
    print(f"    {description}")
    print(f"    current: {shown}")
    print("    hover the mouse over it and press F8  (F9 keeps current, Esc quits)")

    key = wait_for_key([CAPTURE_KEY, SKIP_KEY, QUIT_KEY])
    if key == QUIT_KEY:
        raise Quit()
    if key == SKIP_KEY:
        print(f"    kept {shown}")
        return None

    pos = pyautogui.position()
    print(f"    captured ({pos.x}, {pos.y})")
    return [pos.x, pos.y]


def calibrate_points(cfg):
    print("\n=== click targets ===")
    for name, description in POINTS:
        captured = capture_point(name, description, cfg["points"].get(name))
        if captured:
            cfg["points"][name] = captured


def calibrate_probes(cfg):
    print("\n=== pixel probes ===")
    for name, kind, description in PROBES:
        existing = cfg["probes"].get(name, {})
        captured = capture_point(name, description, existing.get("pos"))
        if not captured:
            continue

        cfg["probes"].setdefault(name, {})["pos"] = captured
        if kind == "color":
            # sample the same block the bot reads, not the cursor
            shot = pyautogui.screenshot(
                region=(captured[0], captured[1], config.PROBE_SIZE, config.PROBE_SIZE)
            )
            colour = list(shot.getpixel((0, 0))[:3])
            cfg["probes"][name]["color"] = colour
            print(f"    colour there is {tuple(colour)}")
            print("    (make sure the portal is actually visible. if that looks")
            print("     like background, press Esc, show the portal and start over)")


def calibrate_bands(cfg):
    print("\n=== vend alignment bands ===")
    print("These are the left and right x-bounds your character has to end up")
    print("between to be able to click a vending machine. Stand at each edge of")
    print("the workable spot and hover over your character.")

    for name, description in BANDS:
        existing = cfg["alignment"].get(name, [None, None])
        left = capture_point(
            f"{name} (left bound)",
            f"the leftmost your character can stand and still reach {description}",
            [existing[0], 0] if existing[0] else None,
        )
        right = capture_point(
            f"{name} (right bound)",
            f"the rightmost your character can stand and still reach {description}",
            [existing[1], 0] if existing[1] else None,
        )

        band = list(existing)
        if left:
            band[0] = left[0]
        if right:
            band[1] = right[0]
        if band[0] is None or band[1] is None:
            print(f"    skipped, {name} still needs both bounds")
            continue
        if band[0] > band[1]:
            band = [band[1], band[0]]
        cfg["alignment"][name] = band
        print(f"    band for {name}: {band[0]}..{band[1]}")


def calibrate_placement(cfg):
    print("\n=== block placement ===")
    print("The put-* bots click one tile to the side of your character. The next")
    print("two captures are measured against each other rather than stored, so")
    print("F9 on either of them leaves x_offset alone.")

    origin = capture_point("character", "the centre of your character", None)
    if origin:
        target = capture_point(
            "placement target", "the tile the bot should click to place a block", None
        )
        if target:
            cfg["placement"]["x_offset"] = target[0] - origin[0]
            print(f"    x_offset: {cfg['placement']['x_offset']}")

    # placing only happens between these two, so nothing overshoots
    lower = capture_point(
        "min_x",
        "the leftmost x where placing is still allowed",
        [cfg["placement"].get("min_x", 0), 0],
    )
    if lower:
        cfg["placement"]["min_x"] = lower[0]

    upper = capture_point(
        "max_x",
        "the rightmost x where placing is still allowed",
        [cfg["placement"].get("max_x", 0), 0],
    )
    if upper:
        cfg["placement"]["max_x"] = upper[0]

    if cfg["placement"].get("min_x", 0) > cfg["placement"].get("max_x", 0):
        cfg["placement"]["min_x"], cfg["placement"]["max_x"] = (
            cfg["placement"]["max_x"],
            cfg["placement"]["min_x"],
        )
        print("    (swapped, min_x has to be the smaller one)")


def backup_asset(path):
    """Keep the original around the first time we overwrite it."""
    if not os.path.exists(path):
        return
    if not os.path.isdir(BACKUP_DIR):
        os.makedirs(BACKUP_DIR)
    backup = os.path.join(BACKUP_DIR, os.path.basename(path))
    if not os.path.exists(backup):
        shutil.copy2(path, backup)


def capture_asset(filename, description):
    print(f"\n  {filename}")
    print(f"    {description}")
    print("    press F8 at the TOP-LEFT corner, then F8 at the BOTTOM-RIGHT")
    print("    (F9 keeps the existing image, Esc quits)")

    key = wait_for_key([CAPTURE_KEY, SKIP_KEY, QUIT_KEY])
    if key == QUIT_KEY:
        raise Quit()
    if key == SKIP_KEY:
        print("    kept existing image")
        return

    first = pyautogui.position()
    print(f"    corner 1: ({first.x}, {first.y}), now the opposite corner")

    key = wait_for_key([CAPTURE_KEY, QUIT_KEY])
    if key == QUIT_KEY:
        raise Quit()
    second = pyautogui.position()

    left, right = sorted((first.x, second.x))
    top, bottom = sorted((first.y, second.y))
    if right - left < 2 or bottom - top < 2:
        print("    that region is too small, skipping")
        return

    # give the cursor time to move out of the crop
    print("    move the mouse off the region, capturing in 2s")
    time.sleep(2)

    crop = pyautogui.screenshot().crop((left, top, right, bottom))
    path = config.asset(filename)
    backup_asset(path)
    crop.save(path)
    print(f"    saved {filename} ({crop.width}x{crop.height})")


def calibrate_assets():
    print("\n=== template images ===")
    print("Each of these is matched pixel-for-pixel, so crop tightly and make")
    print("sure the thing is actually on screen before you capture it.")
    print("Originals are backed up to assets/_backup/ the first time.")

    for filename, description in ASSETS:
        capture_asset(filename, description)


def save(cfg):
    cfg["calibrated_at"]["resolution"] = list(pyautogui.size())

    # pull number arrays back inline, so it stays readable
    text = json.dumps(cfg, indent=2)
    text = re.sub(
        r"\[[\s\d,.-]+?\]",
        lambda m: "[" + " ".join(m.group(0).strip("[]").split()) + "]",
        text,
    )

    with open(config.CONFIG_PATH, "w", encoding="utf-8") as f:
        f.write(text + "\n")
    print(f"\nwrote {config.CONFIG_PATH}")


def main():
    width, height = pyautogui.size()
    print(__doc__)
    print(f"screen: {width}x{height}")
    print("Keys while capturing:  F8 = capture   F9 = keep current   Esc = quit")
    print("Switch to Growtopia now. This window does not need focus, the")
    print("hotkeys are global. Results print back here.")

    with open(config.CONFIG_PATH, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    print("\nWhat do you want to calibrate?")
    print("  1  everything (coordinates, probes, bands, placement, images)")
    print("  2  coordinates, probes, bands and placement only")
    print("  3  template images only")
    print("  4  pixel probes only")
    print("  5  vend alignment bands only")
    print("  6  block placement only")
    choice = input("choice [1]: ").strip() or "1"

    try:
        if choice in ("1", "2"):
            calibrate_points(cfg)
            calibrate_probes(cfg)
            calibrate_bands(cfg)
            calibrate_placement(cfg)
            save(cfg)
        if choice == "4":
            calibrate_probes(cfg)
            save(cfg)
        if choice == "5":
            calibrate_bands(cfg)
            save(cfg)
        if choice == "6":
            calibrate_placement(cfg)
            save(cfg)
        if choice in ("1", "3"):
            calibrate_assets()
    except Quit:
        # quitting part way through still keeps what was captured
        if choice in ("1", "2", "4", "5", "6"):
            save(cfg)
        print("\nstopped early, what you captured up to here has been saved")
        return

    print("\nDone. Run a bot to check it, and alt-tab away to stop it.")


if __name__ == "__main__":
    main()