#!/usr/bin/env python3
"""Synthetic test bench for mk_cal_analyze.py: simulates the B4 galvo (MeerK40t view transform + .cor table +
a made-up physical distortion), burns the calibration pattern onto virtual black paper and 'scans' it at 600 dpi.
Only for verifying the analysis/correction maths. Grok 2026-10-31."""
import argparse, json, math, os, sys
import numpy as np, cv2
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mk_cal_design as D
from mk_view_table import VIEW_MATRIX
from mk_cal_analyze import read_cor

STOCK = None


def interp_tab(tab, t):
    g = np.clip(t / 1024.0, 0, 64 - 1e-9)
    k0, j0 = np.floor(g[:, 0]).astype(int), np.floor(g[:, 1]).astype(int)
    fk, fj = g[:, 0] - k0, g[:, 1] - j0
    v = (tab[j0, k0] * ((1 - fk) * (1 - fj))[:, None] + tab[j0, k0 + 1] * (fk * (1 - fj))[:, None] +
         tab[j0 + 1, k0] * ((1 - fk) * fj)[:, None] + tab[j0 + 1, k0 + 1] * (fk * fj)[:, None])
    return v


class Machine:
    def __init__(self, A_g, true_lens, stock_tab, dist):
        self.A = np.array(A_g, float); self.k0 = 65535.0 / true_lens; self.stock = stock_tab; self.dist = dist

    def phys(self, out):
        u = out - 32768.0
        v = u.copy()
        for _ in range(8):                      # invert stock table: v - stock(C+v) = u
            v = u + interp_tab(self.stock, v + 32768.0)
        g = v / self.k0                         # mm in galvo frame (stock cor makes this ideal)
        g = g + self.dist(g)
        return g @ self.A.T


def default_dist(g):
    x, y = g[:, 0] / 50, g[:, 1] / 50
    dx = 0.30 * x + 0.20 * y + 0.15 * x * y + 0.25 * x * y * y + 0.05 * x * x
    dy = -0.15 * y + 0.10 * x * x - 0.12 * x * y + 0.20 * y * x * x
    return np.c_[dx, dy]


def burn_polys(machine, lens, flags, cor_tab, sign=1.0):
    try:
        import matplotlib.textpath  # noqa: F401
        from gen_mk_cal import text_polys
    except Exception:            # no matplotlib (laptop venv): simulate without the text labels
        text_polys = None
    R = np.array(VIEW_MATRIX[flags], float); k = 65535.0 / lens
    polys = [p for _, p in D.geometry() if _ != 'frame']
    for t, x0, yb, h in (D.text_items('S? MK L131.5') if text_polys else []):
        polys += text_polys(t, x0, yb, h)
    out = []
    for p in polys:
        p = np.array(p, float)
        seg = [p[0:1]]
        for a, b in zip(p[:-1], p[1:]):
            n = max(1, int(np.linalg.norm(b - a) / 0.2))
            seg.append(a + (b - a) * (np.arange(1, n + 1) / n)[:, None])
        p = np.concatenate(seg)
        t = 32768.0 + (p @ R.T) * k
        o = t - sign * interp_tab(cor_tab, t)
        out.append(machine.phys(o))
    return out


def render(polys, out, paper_rot=0.6, paper_off=(108.0, 125.0), scan_rot=180, scan_ang=0.4, seed=1):
    rng = np.random.default_rng(seed)
    ppmx, ppmy = 600 / 25.4, 599.70 / 25.4
    W, Hh = int(215.9 * ppmx), int(297 * ppmy)
    pw, ph = 215.9, 279.4
    def to_paper(q):
        c, s = math.cos(math.radians(paper_rot)), math.sin(math.radians(paper_rot))
        return q @ np.array([[c, -s], [s, c]]).T + np.array(paper_off)
    def to_scan(pp):        # paper coords -> scan mm: paper rotated by scan_rot+scan_ang about its centre, placed on glass
        a = math.radians(scan_rot + scan_ang); c, s = math.cos(a), math.sin(a)
        q = (pp - np.array([pw / 2, ph / 2])) @ np.array([[c, -s], [s, c]]).T
        dims = (pw, ph) if scan_rot % 180 == 0 else (ph, pw)
        return q + np.array([dims[0] / 2 + 0.3, dims[1] / 2 + 3.0])
    img = np.full((Hh, W), 228, np.float32)
    paper = np.array([[8, 0], [pw, 0], [pw, ph], [0, ph], [0, 8]], float)   # snip at paper back-left
    sp = to_scan(paper) * [ppmx, ppmy]
    cv2.fillPoly(img, [np.round(sp * 16).astype(np.int32)], 32.0, cv2.LINE_AA, shift=4)
    lay = np.zeros_like(img)
    for q in polys:
        sq = to_scan(to_paper(q)) * [ppmx, ppmy]
        cv2.polylines(lay, [np.round(sq * 16).astype(np.int32)], False, 1.0, 3, cv2.LINE_AA, shift=4)
    img = img + lay * 95 * (img < 100)
    img = cv2.GaussianBlur(img, (0, 0), 1.0) + rng.normal(0, 4, img.shape).astype(np.float32)
    cv2.imwrite(out, np.clip(img, 0, 255).astype(np.uint8))


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', required=True)
    ap.add_argument('--lens', type=float, default=131.5)
    ap.add_argument('--flags', default='0,1,1,0')
    ap.add_argument('--cor', required=True)
    ap.add_argument('--stock', required=True)
    ap.add_argument('--true-lens', type=float, default=131.5 * 1.012)
    ap.add_argument('--machine-A', default=None, help='json 2x2 galvo->paper; default = inverse of 0,1,1,0 view (correct)')
    ap.add_argument('--board-sign', type=float, default=1.0)
    ap.add_argument('--scan-rot', type=int, default=180)
    ap.add_argument('--seed', type=int, default=1)
    a = ap.parse_args()
    fl = tuple(int(v) for v in a.flags.split(','))
    flags = (bool(fl[0]), bool(fl[1]), bool(fl[2]), fl[3])
    A = np.array(json.loads(a.machine_A)) if a.machine_A else np.array(VIEW_MATRIX[(False, True, True, 0)], float).T
    stock = read_cor(a.stock)[1]
    cor = read_cor(a.cor)[1]
    m = Machine(A, a.true_lens, stock, default_dist)
    render(burn_polys(m, a.lens, flags, cor, a.board_sign), a.out, scan_rot=a.scan_rot, seed=a.seed)
    print('wrote', a.out)