#!/home/sam/dev/perso/flsun/klippy-env/bin/python
"""Migre la bed mesh vers une grille 7x7 (round_probe_count 7, mesh_radius 120).

    ./mesh-to-7.py            # montre le resultat, demande confirmation, ecrit, redemarre

Principe : on ne re-sonde rien. La mesh 5x5 actuelle (rayon 110) definit la
surface que Klipper CROIT aujourd'hui ; on l'echantillonne (interpolation
bilineaire sur la matrice remplie, extrapolation plate au-dela de 110) aux 29
points de la nouvelle grille, on remplit les colonnes hors disque comme
Klipper (copie du bord), et on ecrit :
  - [bed_mesh]            mesh_radius: 120  /  round_probe_count: 7
  - [bed_mesh default]    points 7x7, x_count/y_count 7, min/max +-120
Ensuite : gcodes/first_layer_29_pla.gcode (29 carres de 10 mm) + ./fix-mesh.py
(qui lit la geometrie de la grille dans printer.cfg) pour affiner.
Refuse si une impression est en cours. Sauvegarde de printer.cfg avant ecriture.
"""
import json
import math
import os
import re
import shutil
import socket
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
CFG = os.path.join(HERE, "printer.cfg")
UDS = "/tmp/klippy_uds"
N_NEW, R_NEW = 7, 120.0


def en_impression():
    try:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.settimeout(3)
        s.connect(UDS)
        s.sendall(json.dumps({"id": 1, "method": "objects/query",
                              "params": {"objects": {"print_stats": None}}}).encode() + b"\x03")
        buf = b""
        while b"\x03" not in buf:
            buf += s.recv(65536)
        s.close()
        return json.loads(buf.split(b"\x03")[0])["result"]["status"]["print_stats"]["state"] in ("printing", "paused")
    except OSError:
        return False


def grille(n, r):
    """Reproduit bed_mesh.generate_points() pour un plateau rond : pas, coordonnees, points reels."""
    x_dist = math.floor((2 * r / (n - 1)) * 100) / 100
    new_r = (n // 2) * x_dist
    coords = [round(-new_r + i * x_dist, 3) for i in range(n)]
    reels = {i: [j for j, x in enumerate(coords) if math.hypot(x, coords[i]) <= r] for i in range(n)}
    return x_dist, coords, reels


def lire_mesh(cfg):
    m = re.search(r"^#\*# points =\n((?:#\*# \t.*\n)+)", cfg, re.M)
    if not m:
        sys.exit("bloc '#*# points =' introuvable")
    mat = [[float(v) for v in l.replace("#*# \t", "").split(",")] for l in m.group(1).strip("\n").split("\n")]
    n = len(mat)
    r = float(re.search(r"^mesh_radius:\s*([0-9.]+)", cfg, re.M).group(1))
    return m, mat, n, r


def bilineaire(mat, coords, x, y):
    """Interpolation bilineaire sur la matrice remplie ; plate au-dela du bord."""
    lo, hi = coords[0], coords[-1]
    x = min(max(x, lo), hi)
    y = min(max(y, lo), hi)
    d = coords[1] - coords[0]
    fx, fy = (x - lo) / d, (y - lo) / d
    i, j = min(int(fy), len(coords) - 2), min(int(fx), len(coords) - 2)
    ty, tx = fy - i, fx - j
    return ((1 - ty) * ((1 - tx) * mat[i][j] + tx * mat[i][j + 1])
            + ty * ((1 - tx) * mat[i + 1][j] + tx * mat[i + 1][j + 1]))


def main():
    print(__doc__)
    if en_impression():
        sys.exit("impression en cours : on ne touche pas a la mesh maintenant")
    cfg = open(CFG).read()
    m, mat, n_old, r_old = lire_mesh(cfg)
    _, coords_old, _ = grille(n_old, r_old)
    x_dist, coords, reels = grille(N_NEW, R_NEW)
    print("ancienne grille : %dx%d, rayon %g  ->  nouvelle : %dx%d, rayon %g, pas %g, %d points reels"
          % (n_old, n_old, r_old, N_NEW, N_NEW, R_NEW, x_dist, sum(len(v) for v in reels.values())))

    nouveau = [[0.0] * N_NEW for _ in range(N_NEW)]
    for i in range(N_NEW):
        for j in reels[i]:
            nouveau[i][j] = bilineaire(mat, coords_old, coords[j], coords[i])
        lo, hi = reels[i][0], reels[i][-1]
        for j in range(N_NEW):                      # remplissage hors disque = copie du bord (comme Klipper)
            if j < lo:
                nouveau[i][j] = nouveau[i][lo]
            elif j > hi:
                nouveau[i][j] = nouveau[i][hi]

    print("\nnouvelle matrice (lignes Y %g -> %g) :" % (coords[0], coords[-1]))
    for row in nouveau:
        print("  " + ", ".join("%7.3f" % v for v in row))

    if input("\nEcrire (config + profil) et redemarrer klippy (o/N) ? ").strip().lower() != "o":
        sys.exit("rien ecrit")
    shutil.copy(CFG, CFG.replace(".cfg", "-avant-mesh7-%s.cfg" % time.strftime("%Y%m%d_%H%M%S")))

    # 1) section [bed_mesh] du corps
    cfg = re.sub(r"^mesh_radius:\s*[0-9.]+", "mesh_radius: %g" % R_NEW, cfg, count=1, flags=re.M)
    cfg = re.sub(r"^round_probe_count:\s*\d+", "round_probe_count: %d" % N_NEW, cfg, count=1, flags=re.M)
    # 2) profil sauvegarde
    bloc = "#*# points =\n" + "".join("#*# \t  " + ", ".join("%.6f" % v for v in row) + "\n" for row in nouveau)
    cfg = cfg[:m.start()] + bloc + cfg[m.end():]
    for cle, val in (("x_count", N_NEW), ("y_count", N_NEW),
                     ("min_x", -R_NEW), ("max_x", R_NEW), ("min_y", -R_NEW), ("max_y", R_NEW)):
        cfg = re.sub(r"^(#\*# %s = ).*$" % cle, lambda mm, v=val: mm.group(1) + ("%g" % v), cfg, count=1, flags=re.M)
    open(CFG, "w").write(cfg)
    print("ecrit.")
    subprocess.run([os.path.join(HERE, "start-klipper.sh"), "restart"])
    print("klippy redemarre : attendre ~15 s, G28, puis first_layer_29_pla + fix-mesh.py")


if __name__ == "__main__":
    main()