#!/usr/bin/env catnip
# Bras robotique à rotors — cinématique directe (algèbre géométrique)
#
# Chaque articulation applique un rotor R = s + B (scalaire + bivecteur). Le
# bivecteur B encode le PLAN de rotation ; son dual est l'axe, et B = sin(θ/2)·axe,
# s = cos(θ/2) donne l'isomorphisme rotor ↔ quaternion unitaire. Le sandwich
# v' = R v R̃ se réduit alors à v' = v + 2s(B×v) + 2 B×(B×v).
#
# Cinématique directe : la rotation cumulée à l'articulation i est le PRODUIT des
# rotors des articulations précédentes. Composer deux rotors = leur produit
# géométrique (un produit de quaternions, ~16 mults), pas un produit matriciel
# 3x3 (~27 mults). Le produit de rotors reste sur la variété des unitaires, donc
# pas de dérive hors SO(3) (aucune ré-orthonormalisation), là où l'enchaînement de
# matrices accumule l'erreur numérique. On vérifie l'équivalence numérique en
# comparant chaque position d'articulation à une chaîne de matrices de Rodrigues.
#
# DEPS: numpy, pillow

numpy = import('numpy')
math = import('math')
Image = import('PIL.Image')
ImageDraw = import('PIL.ImageDraw')
import('pathlib', 'Path')

script_dir = Path(META.file).parent
output_dir = script_dir / 'output'
output_dir.mkdir(exist_ok=True)

struct Vec3 {
    x: float; y: float; z: float;

    arr(self) => { numpy.array(list(self.x, self.y, self.z)) }
}

from_arr = (a): Vec3 => { Vec3(float(a[0]), float(a[1]), float(a[2])) }

# Rotor = partie scalaire + bivecteur (dual de l'axe de rotation).
struct Rotor {
    s: float; b: Vec3;

    # Produit sandwich R v R̃, forme fermée avec deux produits croisés.
    apply(self, v: Vec3): Vec3 => {
        bv = self.b.arr()
        t = 2.0 * numpy.cross(bv, v.arr())
        from_arr(v.arr() + self.s * t + numpy.cross(bv, t))
    }
}

rotor_from_axis_angle = (axis: Vec3, angle: float): Rotor => {
    k = axis.arr() / numpy.linalg.norm(axis.arr())
    half = angle / 2.0
    Rotor(math.cos(half), from_arr(math.sin(half) * k))
}

# Produit géométrique de deux rotors (produit de quaternions de Hamilton).
# rotor_product(a, b) applique d'abord b puis a : (s_a+B_a)(s_b+B_b). C'est ce
# qui remplace le produit de matrices de rotation dans la chaîne cinématique.
rotor_product = (a: Rotor, b: Rotor): Rotor => {
    ua = a.b.arr()
    ub = b.b.arr()
    s = a.s * b.s - float(numpy.dot(ua, ub))
    v = a.s * ub + b.s * ua + numpy.cross(ua, ub)
    Rotor(s, from_arr(v))
}

# Matrice de rotation de référence (Rodrigues), chemin indépendant du rotor.
rotation_matrix = (axis: Vec3, angle: float) => {
    k = axis.arr() / numpy.linalg.norm(axis.arr())
    cross_k = numpy.array(list(
            list(0.0, -k[2], k[1]),
            list(k[2], 0.0, -k[0]),
            list(-k[1], k[0], 0.0),
        ))
    numpy.eye(3) * math.cos(angle) + math.sin(angle) * cross_k + (1.0 - math.cos(angle)) * numpy.outer(k, k)
}

# Une articulation : axe de rotation local + longueur du segment attaché (le
# segment part le long de +x dans le repère de l'articulation).
struct Joint {
    axis: Vec3; length: float;
}

struct Arm {
    joints: list[Joint];

    # Cinématique directe par rotors. Retourne les positions (base + chaque
    # extrémité de segment). q est la rotation cumulée : q_i = q_{i-1} · R_i,
    # obtenue par produit de rotors, appliquée au segment local.
    fk(self, angles: list[float]): list[Vec3] => {
        positions = list(Vec3(0.0, 0.0, 0.0))
        q = Rotor(1.0, Vec3(0.0, 0.0, 0.0))
        p = Vec3(0.0, 0.0, 0.0)
        i = 0
        while i < len(self.joints) {
            joint = self.joints[i]
            q = rotor_product(q, rotor_from_axis_angle(joint.axis, angles[i]))
            tip = q.apply(Vec3(joint.length, 0.0, 0.0))
            p = from_arr(p.arr() + tip.arr())
            positions.append(p)
            i = i + 1
        }
        positions
    }
}

# Cinématique directe de référence par matrices : m_i = m_{i-1} @ R_i, la
# position de segment est m_i @ (length, 0, 0). Chaîne entièrement indépendante
# du rotor, sert de contrôle numérique.
fk_matrix = (arm: Arm, angles: list[float]) => {
    positions = list(numpy.array(list(0.0, 0.0, 0.0)))
    m = numpy.eye(3)
    p = numpy.array(list(0.0, 0.0, 0.0))
    i = 0
    while i < len(arm.joints) {
        joint = arm.joints[i]
        m = m.dot(rotation_matrix(joint.axis, angles[i]))
        p = p + m.dot(numpy.array(list(joint.length, 0.0, 0.0)))
        positions.append(p)
        i = i + 1
    }
    positions
}

struct Pose {
    name: str; angles: list[float];
}

AXIS_Z = Vec3(0.0, 0.0, 1.0)
AXIS_Y = Vec3(0.0, 1.0, 0.0)

# Bras à trois articulations : lacet à la base (axe z), puis deux tangages
# (axe y) pour l'épaule et le coude. Segments de plus en plus courts.
arm = Arm(list(
        Joint(AXIS_Z, 1.6),
        Joint(AXIS_Y, 1.3),
        Joint(AXIS_Y, 1.0),
    ))

# Presets de pose : chaque pose fixe un angle par articulation (contrainte
# implicite : autant d'angles que d'articulations).
poses = list(
    Pose('home', list(0.0, 0.0, 0.0)),
    Pose('reach', list(0.0, -0.9, -0.6)),
    Pose('wave', list(0.7, -0.5, 0.9)),
    Pose('fold', list(0.2, 1.3, 1.5)),
)

# --- Rendu filaire des poses ---------------------------------------------------

WIDTH = 900
HEIGHT = 980
TITLE_H = 60
PANEL_W = WIDTH / 2
PANEL_H = (HEIGHT - TITLE_H) / 2
FOCAL = 560.0
CAMDIST = 6.5
NEAR = 0.1

# Vue 3/4 fixe, elle-même une composition de deux rotors (lacet puis tangage) :
# la même machinerie de composition sert à la caméra et au bras.
view = rotor_product(
    rotor_from_axis_angle(Vec3(1.0, 0.0, 0.0), -0.45),
    rotor_from_axis_angle(AXIS_Y, 0.6),
)

seg_colors = list(
    tuple(230, 90, 90),
    tuple(90, 200, 120),
    tuple(90, 150, 240),
)

project = (v: Vec3, cx: float, cy: float) => {
    depth = CAMDIST - v.z
    depth2 = if depth > NEAR { depth } else { NEAR }
    tuple(cx + FOCAL * v.x / depth2, cy - FOCAL * v.y / depth2)
}

image = Image.new('RGB', tuple(WIDTH, HEIGHT), tuple(18, 18, 24))
draw = ImageDraw.Draw(image)
draw.text(tuple(18, 20), "Bras robotique a rotors - cinematique directe", tuple(235, 235, 240))

disc = (center, radius: float, fill, outline) => {
    x = center[0]
    y = center[1]
    draw.ellipse(list(tuple(x - radius, y - radius), tuple(x + radius, y + radius)), fill=fill, outline=outline)
}

# Dessine une pose dans son panneau et renvoie l'erreur max rotor-vs-matrice.
draw_pose = (idx: int, pose: Pose): float => {
    col = idx % 2
    row = int(idx / 2)
    ox = col * PANEL_W
    oy = TITLE_H + row * PANEL_H
    cx = ox + PANEL_W / 2.0
    cy = oy + PANEL_H / 2.0
    draw.rectangle(list(tuple(ox + 1, oy + 1), tuple(ox + PANEL_W - 2, oy + PANEL_H - 2)), outline=tuple(45, 45, 58))

    rp = arm.fk(pose.angles)

    # Contrôle numérique : mêmes angles, chaîne de matrices de Rodrigues.
    ref = numpy.array(fk_matrix(arm, pose.angles))
    got = numpy.array(rp.[(v) => { v.arr() }])
    err = float(numpy.max(numpy.linalg.norm(ref - got, axis=1)))

    # Vue 3/4 puis recentrage sur le barycentre pour cadrer chaque pose.
    viewed = rp.[(v) => { view.apply(v) }]
    centroid = reduce(viewed.[(v) => { v.arr() }], (a, b) => { a + b }) / len(viewed)
    projected = viewed.[(v) => { project(from_arr(v.arr() - centroid), cx, cy) }]

    # Segments filaires reliant les articulations.
    i = 0
    while i < len(projected) - 1 {
        draw.line(list(projected[i], projected[i + 1]), fill=seg_colors[i % len(seg_colors)], width=9)
        i = i + 1
    }

    # Articulations (noeuds), base carrée, effecteur mis en évidence.
    last = len(projected) - 1
    j = 0
    while j <= last {
        pt = projected[j]
        if j == 0 {
            disc(pt, 9.0, tuple(210, 210, 220), tuple(20, 20, 26))
        } else {
            if j == last {
                disc(pt, 11.0, tuple(250, 210, 70), tuple(20, 20, 26))
            } else {
                disc(pt, 7.0, tuple(60, 60, 72), tuple(210, 210, 220))
            }
        }
        j = j + 1
    }

    eff = rp[last]
    label = f"{pose.name}  eff=({round(eff.x, 2)}, {round(eff.y, 2)}, {round(eff.z, 2)})"
    draw.text(tuple(ox + 14, oy + 12), label, tuple(225, 225, 232))
    err
}

print("⇒ Cinématique directe par rotors (contrôle vs matrices de Rodrigues)")
errors = list()
k = 0
while k < len(poses) {
    pose = poses[k]
    err = draw_pose(k, pose)
    errors.append(err)
    eff = arm.fk(pose.angles)[len(arm.joints)]
    print(f"  {pose.name:>6} : effecteur=({round(eff.x, 3)}, {round(eff.y, 3)}, {round(eff.z, 3)})  erreur={err}")
    k = k + 1
}

max_error = float(numpy.max(numpy.array(errors)))
print(f"  contrôle rotor vs matrice : erreur max (toutes poses) = {max_error}")

output_path = output_dir / 'rotor_robot_arm.png'
image.save(str(output_path))
print()
print(f"⇒ Image → {output_path}")

# Aperçu navigateur : sert l'image une fois puis rend la main.
# --no-browser garde un chemin headless (le PNG reste écrit ci-dessus).
if '--no-browser' not in import('sys').argv {
    http = import('http')
    b64 = import('base64').b64encode(output_path.read_bytes()).decode('ascii')
    http.serve(f'<!doctype html><meta charset="utf-8"><title>Rotor robot arm</title><body style="margin:0;background:#0d1117"><img style="max-width:100%;display:block;margin:0 auto" src="data:image/png;base64,{b64}">',
        0, 'text/html; charset=utf-8', True)
}