#!/usr/bin/env python3
"""On-screen camera calibration: ChArUco board on a second monitor, live capture,
KB4 (Kannala-Brandt 'kb4') + pinhole-radtan calibration.

Pick which monitor shows the pattern with --screen; its geometry (resolution,
position, physical size) is read from `xrandr` automatically. List screens first:

    python3 calibrate_camera.py --list-screens
    python3 calibrate_camera.py --screen HDMI-0 --device /dev/video0 --width 1280 --height 800

If xrandr is unavailable (headless/Wayland) or reports the wrong physical size,
set the pattern-monitor geometry manually instead (these override --screen):

    python3 calibrate_camera.py \
        --device /dev/video0 --width 1280 --height 800 \
        --monitor-x 1920 --monitor-px-w 2560 --monitor-px-h 1440 --monitor-mm-w 697

Workflow:
  1. The pattern window opens fullscreen on the second monitor; the live camera
     preview opens on the primary monitor.
  2. Point the camera at the pattern from 1.5-2.5 m (lens is focused at infinity;
     farther = sharper corners). Fill different parts of the image, tilt the
     camera up to ~35 deg in all directions, and get the board into the corners
     of the FOV. Samples auto-capture when the view is sharp, well-detected, and
     different from previous samples.
  3. Collect 30+ samples with good coverage, press 'c' to calibrate.

Keys (camera window focused):
  SPACE  force-capture current frame
  c      run calibration with collected samples
  u      drop the last sample
  r      reset all samples
  q/ESC  quit

Re-run calibration offline from saved frames:
    python3 calibrate_camera.py --from-dir calibration_results/<ts>
"""

import argparse
import json
import math
import re
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np

# ---------------------------------------------------------------------------
# Monitors / screen selection
# ---------------------------------------------------------------------------

# default pattern-monitor geometry, used when neither --screen nor --monitor-*
# flags are given (px width, px height, x offset, physical width in mm)
DEFAULT_GEOMETRY = {"monitor_px_w": 2560, "monitor_px_h": 1440,
                    "monitor_x": 1920, "monitor_mm_w": 697.0}

_XRANDR_MON = re.compile(
    r"^\s*(\d+):\s+\+(\*?)(\S+)\s+(\d+)/(\d+)x(\d+)/(\d+)\+(\d+)\+(\d+)")


def list_monitors():
    """Parse `xrandr --listmonitors` into monitor dicts (X11 only; [] otherwise)."""
    try:
        out = subprocess.run(["xrandr", "--listmonitors"],
                             capture_output=True, text=True, check=True).stdout
    except (FileNotFoundError, subprocess.CalledProcessError):
        return []
    monitors = []
    for line in out.splitlines():
        m = _XRANDR_MON.match(line)
        if not m:
            continue
        idx, primary, name, pw, mw, ph, mh, x, y = m.groups()
        monitors.append({
            "index": int(idx), "name": name, "primary": bool(primary),
            "px_w": int(pw), "px_h": int(ph),
            "mm_w": float(mw), "mm_h": float(mh),
            "x": int(x), "y": int(y),
        })
    return monitors


def resolve_screen(token, monitors):
    """Find a monitor by xrandr index or output name (exact, then substring)."""
    if token.isdigit():
        for m in monitors:
            if m["index"] == int(token):
                return m
    for m in monitors:
        if m["name"] == token:
            return m
    for m in monitors:
        if token.lower() in m["name"].lower():
            return m
    return None


def describe_monitor(m):
    star = " *primary" if m["primary"] else ""
    return (f"{m['index']}: {m['name']:<12s} {m['px_w']}x{m['px_h']} "
            f"@ +{m['x']}+{m['y']}  {m['mm_w']:.0f}x{m['mm_h']:.0f} mm{star}")


def resolve_geometry(args):
    """Fill args.monitor_* from --screen (xrandr) and/or explicit overrides."""
    monitors = list_monitors()
    geo = dict(DEFAULT_GEOMETRY)

    if args.screen is not None:
        if not monitors:
            sys.exit("--screen requested but no monitors detected (needs xrandr / "
                     "X11); pass --monitor-x/-px-w/-px-h/-mm-w manually instead")
        mon = resolve_screen(args.screen, monitors)
        if mon is None:
            avail = ", ".join(f"{m['index']}:{m['name']}" for m in monitors)
            sys.exit(f"screen '{args.screen}' not found; available: {avail}")
        if mon["mm_w"] <= 0 and args.monitor_mm_w is None:
            sys.exit(f"screen '{mon['name']}' reports 0 mm physical width; "
                     "pass --monitor-mm-w <mm> explicitly")
        geo.update(monitor_x=mon["x"], monitor_px_w=mon["px_w"],
                   monitor_px_h=mon["px_h"], monitor_mm_w=float(mon["mm_w"]))
        print(f"using screen {mon['index']}:{mon['name']} {mon['px_w']}x{mon['px_h']} "
              f"@ +{mon['x']} ({mon['mm_w']:.0f} mm wide)")
    elif len(monitors) > 1:
        print("tip: multiple screens detected; pick one with --list-screens / --screen")

    # explicit CLI flags win over both defaults and --screen
    for key in ("monitor_x", "monitor_px_w", "monitor_px_h", "monitor_mm_w"):
        override = getattr(args, key)
        if override is not None:
            geo[key] = override
        setattr(args, key, geo[key])


# ---------------------------------------------------------------------------
# Board / pattern
# ---------------------------------------------------------------------------


def make_board(args):
    dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_100)
    px_per_mm = args.monitor_px_w / args.monitor_mm_w
    margin_px = 60
    square_px = min((args.monitor_px_w - 2 * margin_px) // args.squares_x,
                    (args.monitor_px_h - 2 * margin_px) // args.squares_y)
    square_mm = square_px / px_per_mm
    square_m = square_mm * 1e-3
    board = cv2.aruco.CharucoBoard((args.squares_x, args.squares_y),
                                   square_m, square_m * 0.75, dictionary)
    img_w = square_px * args.squares_x + 2 * margin_px
    img_h = square_px * args.squares_y + 2 * margin_px
    board_img = board.generateImage((img_w, img_h), marginSize=margin_px)
    canvas = np.full((args.monitor_px_h, args.monitor_px_w), 255, np.uint8)
    y0 = (args.monitor_px_h - img_h) // 2
    x0 = (args.monitor_px_w - img_w) // 2
    canvas[y0:y0 + img_h, x0:x0 + img_w] = board_img
    print(f"board: {args.squares_x}x{args.squares_y} squares, "
          f"{square_px}px = {square_mm:.2f}mm per square")
    return board, canvas


# ---------------------------------------------------------------------------
# Camera
# ---------------------------------------------------------------------------


def open_camera(args):
    cap = cv2.VideoCapture(args.device, cv2.CAP_V4L2)
    if not cap.isOpened():
        sys.exit(f"cannot open {args.device}")
    for fourcc in (args.fourcc, "MJPG", "GREY", "YUYV"):
        if not fourcc:
            continue
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*fourcc))
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.height)
        ok, frame = cap.read()
        if ok and frame is not None:
            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            print(f"camera: {args.device} {w}x{h} fourcc={fourcc} "
                  f"fps={cap.get(cv2.CAP_PROP_FPS):.0f}")
            if (w, h) != (args.width, args.height):
                print(f"WARNING: negotiated {w}x{h} != requested "
                      f"{args.width}x{args.height} - calibration applies to {w}x{h}")
            return cap
    sys.exit("camera opened but no frames received with any pixel format")


def to_gray(frame):
    return frame if frame.ndim == 2 else cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)


# ---------------------------------------------------------------------------
# Calibration
# ---------------------------------------------------------------------------


def calibrate(samples, image_size, board):
    """samples: list of (charuco_corners Nx1x2 float32, charuco_ids Nx1 int32)."""
    all_obj = board.getChessboardCorners().astype(np.float64)  # (N,3)
    obj_pts, img_pts = [], []
    for corners, ids in samples:
        ids = ids.ravel()
        if len(ids) < 8:
            continue
        obj_pts.append(all_obj[ids].reshape(-1, 1, 3))
        img_pts.append(corners.astype(np.float64).reshape(-1, 1, 2))
    if len(obj_pts) < 6:
        print(f"not enough usable views ({len(obj_pts)}), need >= 6")
        return None

    out = {"views_used": len(obj_pts), "image_size": list(image_size)}

    # --- pinhole + radtan (robust reference) ---
    rms, K, dist, _, _ = cv2.calibrateCamera(
        [o.astype(np.float32) for o in obj_pts],
        [i.astype(np.float32) for i in img_pts],
        image_size, None, None)
    out["pinhole_radtan"] = {
        "rms_px": float(rms),
        "fx": float(K[0, 0]), "fy": float(K[1, 1]),
        "cx": float(K[0, 2]), "cy": float(K[1, 2]),
        "dist": [float(d) for d in dist.ravel()],
    }

    # --- KB4 fisheye with bad-view pruning ---
    flags = (cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC |
             cv2.fisheye.CALIB_FIX_SKEW |
             cv2.fisheye.CALIB_CHECK_COND)
    obj_f = [o.reshape(1, -1, 3) for o in obj_pts]
    img_f = [i.reshape(1, -1, 2) for i in img_pts]
    kb4 = None
    for _ in range(15):
        if len(obj_f) < 6:
            break
        try:
            Kf = np.eye(3)
            Df = np.zeros((4, 1))
            rms_f, Kf, Df, _, _ = cv2.fisheye.calibrate(
                obj_f, img_f, image_size, Kf, Df, flags=flags,
                criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 1e-8))
            ks = np.asarray(Df).ravel()
            kb4 = {
                "rms_px": float(rms_f),
                "fx": float(Kf[0, 0]), "fy": float(Kf[1, 1]),
                "cx": float(Kf[0, 2]), "cy": float(Kf[1, 2]),
                "k1": float(ks[0]), "k2": float(ks[1]),
                "k3": float(ks[2]), "k4": float(ks[3]),
                "views_used": len(obj_f),
            }
            break
        except cv2.error as e:
            msg = str(e)
            idx = None
            if "input array " in msg:  # CALIB_CHECK_COND names the offending view
                try:
                    idx = int(msg.split("input array ")[1].split()[0])
                except (ValueError, IndexError):
                    idx = None
            if idx is None or idx >= len(obj_f):
                # fall back: drop CHECK_COND and accept the fit
                flags &= ~cv2.fisheye.CALIB_CHECK_COND
                continue
            del obj_f[idx], img_f[idx]
    out["kb4"] = kb4
    return out


def report(result):
    print("\n" + "=" * 64)
    ph = result["pinhole_radtan"]
    w, h = result["image_size"]
    print(f"views: {result['views_used']}  image: {w}x{h}")
    print(f"pinhole+radtan: rms={ph['rms_px']:.3f}px  "
          f"fx={ph['fx']:.2f} fy={ph['fy']:.2f} cx={ph['cx']:.1f} cy={ph['cy']:.1f}")
    if result.get("kb4"):
        kb = result["kb4"]
        print(f"kb4:            rms={kb['rms_px']:.3f}px  "
              f"fx={kb['fx']:.2f} fy={kb['fy']:.2f} cx={kb['cx']:.1f} cy={kb['cy']:.1f}")
        print(f"                k1={kb['k1']:.6f} k2={kb['k2']:.6f} "
              f"k3={kb['k3']:.6f} k4={kb['k4']:.6f}")
        fx = kb["fx"]
    else:
        print("kb4 fit failed - using pinhole fx for HFOV estimate")
        fx = ph["fx"]
    hfov = 2 * math.degrees(math.atan((w / 2) / fx))
    print(f"HFOV ~= {hfov:.1f} deg")
    print("=" * 64)


def save_results(out_dir, result, board_meta):
    out_dir.mkdir(parents=True, exist_ok=True)
    result = dict(result, board=board_meta,
                  timestamp=datetime.now().isoformat(timespec="seconds"))
    with open(out_dir / "calibration.json", "w") as f:
        json.dump(result, f, indent=2)
    if result.get("kb4"):
        kb = result["kb4"]
        w, h = result["image_size"]
        yaml_snippet = (
            "# KB4 intrinsics measured with calibrate_camera_onscreen.py\n"
            f"# {result['timestamp']}  rms={kb['rms_px']:.3f}px  views={kb['views_used']}\n"
            "intrinsics:\n"
            f"  fx: {kb['fx']:.4f}\n"
            f"  fy: {kb['fy']:.4f}\n"
            f"  cx: {kb['cx']:.4f}\n"
            f"  cy: {kb['cy']:.4f}\n"
            f"  k1: {kb['k1']:.8f}\n"
            f"  k2: {kb['k2']:.8f}\n"
            f"  k3: {kb['k3']:.8f}\n"
            f"  k4: {kb['k4']:.8f}\n"
            "resolution:\n"
            f"  width: {w}\n"
            f"  height: {h}\n")
        (out_dir / "camera_params_snippet.yaml").write_text(yaml_snippet)
    print(f"saved -> {out_dir}/calibration.json")


# ---------------------------------------------------------------------------
# Live capture loop
# ---------------------------------------------------------------------------

GRID_X, GRID_Y = 6, 4


def coverage_cells(corners, w, h):
    cells = set()
    for c in corners.reshape(-1, 2):
        cells.add((min(GRID_X - 1, int(c[0] / w * GRID_X)),
                   min(GRID_Y - 1, int(c[1] / h * GRID_Y))))
    return cells


def run_live(args):
    board, pattern_img = make_board(args)
    detector = cv2.aruco.CharucoDetector(board)

    # order matters for most window managers: show, pump events, move onto the
    # target monitor, pump events, then go fullscreen
    cv2.namedWindow("CALIBRATION PATTERN", cv2.WINDOW_NORMAL)
    cv2.imshow("CALIBRATION PATTERN", pattern_img)
    cv2.waitKey(100)
    cv2.moveWindow("CALIBRATION PATTERN", args.monitor_x + 100, 100)
    cv2.waitKey(100)
    cv2.setWindowProperty("CALIBRATION PATTERN", cv2.WND_PROP_FULLSCREEN,
                          cv2.WINDOW_FULLSCREEN)
    cv2.waitKey(100)
    cv2.namedWindow("camera", cv2.WINDOW_NORMAL)
    cv2.moveWindow("camera", 100, 100)

    cap = open_camera(args)
    out_dir = (Path(__file__).parent / "calibration_results" /
               datetime.now().strftime("%Y%m%d_%H%M%S"))
    frames_dir = out_dir / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    samples = []          # (charuco_corners, charuco_ids)
    sample_centroids = []
    covered = set()
    last_capture_t = 0.0
    last_result = None
    image_size = None
    board_meta = {"squares_x": args.squares_x, "squares_y": args.squares_y}

    print("\ncollect 30+ samples (tilt the camera, cover all FOV regions), "
          "then press 'c'\n")
    while True:
        ok, frame = cap.read()
        if not ok:
            time.sleep(0.05)
            continue
        gray = to_gray(frame)
        h, w = gray.shape
        image_size = (w, h)
        vis = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

        ch_corners, ch_ids, mk_corners, _ = detector.detectBoard(gray)
        n_det = 0 if ch_ids is None else len(ch_ids)
        sharp = cv2.Laplacian(gray, cv2.CV_64F).var()

        auto_ok = False
        if n_det >= args.min_corners and sharp >= args.min_sharpness:
            cv2.aruco.drawDetectedCornersCharuco(vis, ch_corners, ch_ids,
                                                 (0, 255, 0))
            centroid = ch_corners.reshape(-1, 2).mean(axis=0)
            moved = (not sample_centroids or
                     min(np.linalg.norm(centroid - c) for c in sample_centroids)
                     > args.min_move_px)
            if moved and time.time() - last_capture_t > 0.4:
                auto_ok = True

        key = cv2.waitKey(1) & 0xFF
        force = key == ord(' ')
        if (auto_ok or (force and n_det >= 8)):
            samples.append((ch_corners.copy(), ch_ids.copy()))
            sample_centroids.append(ch_corners.reshape(-1, 2).mean(axis=0))
            covered |= coverage_cells(ch_corners, w, h)
            last_capture_t = time.time()
            cv2.imwrite(str(frames_dir / f"{len(samples):03d}.png"), gray)
            print(f"sample {len(samples):3d}: corners={n_det} sharp={sharp:.0f} "
                  f"coverage={len(covered)}/{GRID_X*GRID_Y}")

        # HUD
        for gx in range(GRID_X):
            for gy in range(GRID_Y):
                color = (0, 160, 0) if (gx, gy) in covered else (0, 0, 180)
                cv2.rectangle(vis, (int(gx * w / GRID_X), int(gy * h / GRID_Y)),
                              (int((gx + 1) * w / GRID_X) - 2,
                               int((gy + 1) * h / GRID_Y) - 2), color, 1)
        flash = (time.time() - last_capture_t) < 0.15
        hud = [
            f"samples {len(samples)}  coverage {len(covered)}/{GRID_X*GRID_Y}  "
            f"corners {n_det}  sharp {sharp:.0f}"
            + ("  [CAPTURE]" if flash else ""),
            "SPACE capture | c calibrate | u undo | r reset | q quit",
        ]
        if last_result and last_result.get("kb4"):
            kb = last_result["kb4"]
            hud.append(f"kb4 fx={kb['fx']:.1f} fy={kb['fy']:.1f} "
                       f"rms={kb['rms_px']:.2f}px")
        for i, line in enumerate(hud):
            cv2.putText(vis, line, (10, 28 + 26 * i), cv2.FONT_HERSHEY_SIMPLEX,
                        0.62, (0, 0, 0), 3, cv2.LINE_AA)
            cv2.putText(vis, line, (10, 28 + 26 * i), cv2.FONT_HERSHEY_SIMPLEX,
                        0.62, (60, 220, 255), 1, cv2.LINE_AA)
        cv2.imshow("camera", vis)

        if key in (ord('q'), 27):
            break
        elif key == ord('u') and samples:
            samples.pop()
            sample_centroids.pop()
            print(f"dropped last sample ({len(samples)} left)")
        elif key == ord('r'):
            samples.clear()
            sample_centroids.clear()
            covered.clear()
            print("samples reset")
        elif key == ord('c') and samples:
            print(f"\ncalibrating with {len(samples)} samples...")
            result = calibrate(samples, image_size, board)
            if result:
                report(result)
                save_results(out_dir, result, board_meta)
                last_result = result

    cap.release()
    cv2.destroyAllWindows()


def run_from_dir(args):
    board, _ = make_board(args)
    detector = cv2.aruco.CharucoDetector(board)
    frames = sorted(Path(args.from_dir).glob("frames/*.png"))
    if not frames:
        sys.exit(f"no frames in {args.from_dir}/frames")
    samples, image_size = [], None
    for p in frames:
        gray = cv2.imread(str(p), cv2.IMREAD_GRAYSCALE)
        image_size = (gray.shape[1], gray.shape[0])
        ch_corners, ch_ids, _, _ = detector.detectBoard(gray)
        if ch_ids is not None and len(ch_ids) >= 8:
            samples.append((ch_corners, ch_ids))
    print(f"re-detected {len(samples)}/{len(frames)} frames")
    result = calibrate(samples, image_size, board)
    if result:
        report(result)
        save_results(Path(args.from_dir), result,
                     {"squares_x": args.squares_x, "squares_y": args.squares_y})


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--device", default="/dev/video0")
    p.add_argument("--width", type=int, default=1280)
    p.add_argument("--height", type=int, default=800)
    p.add_argument("--fourcc", default="")
    # pattern monitor: pick by --screen (auto-read from xrandr), or set geometry
    # manually with the --monitor-* overrides below
    p.add_argument("--screen", default=None,
                   help="monitor to show the pattern on, by xrandr index or name "
                        "(e.g. --screen 1 or --screen HDMI-0); fills geometry")
    p.add_argument("--list-screens", action="store_true",
                   help="list monitors detected via xrandr and exit")
    p.add_argument("--monitor-x", type=int, default=None,
                   help="override: x offset of the pattern monitor (px)")
    p.add_argument("--monitor-px-w", type=int, default=None,
                   help="override: pattern monitor width (px)")
    p.add_argument("--monitor-px-h", type=int, default=None,
                   help="override: pattern monitor height (px)")
    p.add_argument("--monitor-mm-w", type=float, default=None,
                   help="override: physical width of the pattern monitor (mm)")
    p.add_argument("--squares-x", type=int, default=10)
    p.add_argument("--squares-y", type=int, default=7)
    p.add_argument("--min-corners", type=int, default=15,
                   help="min charuco corners for auto-capture")
    p.add_argument("--min-sharpness", type=float, default=60.0,
                   help="min Laplacian variance for auto-capture")
    p.add_argument("--min-move-px", type=float, default=40.0,
                   help="min board centroid movement between auto-captures")
    p.add_argument("--from-dir", default=None,
                   help="recalibrate from saved frames instead of live capture")
    args = p.parse_args()

    if args.list_screens:
        monitors = list_monitors()
        if not monitors:
            sys.exit("no monitors detected (xrandr unavailable or returned nothing; "
                     "X11 only)")
        print("detected screens (use --screen <index|name>):")
        for m in monitors:
            print("  " + describe_monitor(m))
        return

    resolve_geometry(args)

    if args.from_dir:
        run_from_dir(args)
    else:
        run_live(args)


if __name__ == "__main__":
    main()