#!/usr/bin/env python3
"""
4_yAxis_rotation_reconstruction.py
====================================
Berechnet die Rotationsachse eines Gelenks aus drei Board-Erkennungs-Timestamps.

Gegeben drei Messungen (Pos A, B, C) – in denen dieselben fremd-Marker
(link != 'Board') erkannt wurden – bestimmt das Skript:
  - Richtung der Rotationsachse  (axisDir, Einheitsvektor)
  - Referenzpunkt auf der Achse  (axisPoint_mm)
  - Residuen pro Punkt           (Qualitätsmass)

Methode (doc/Kalibrierung.md → [4] Arm1, Verfahren B):
  Jeder Marker bewegt sich auf einem Kreisbogen. Die Ebene des Kreises
  steht senkrecht zur Rotationsachse → Normalenvektor = Achsenrichtung.
  Der Umkreismittelpunkt des Dreiecks P1-P2-P3 liegt auf der Achse.

  Statt nur dem Marker-Zentrum werden alle vier Ecken (corners_m) verwendet:
  das liefert 4× mehr Datenpunkte und macht die Schätzung robuster.

Usage:
  python3 4_yAxis_rotation_reconstruction.py <posA> <posB> <posC> [Optionen]

  <posA/B/C>  Verzeichnis mit aruco_marker_poses.json  ODER  direkter Dateipfad

Optionen:
  --output, -o  Ergebnis in JSON-Datei speichern
  --pretty      Eingerücktes JSON auf stdout
  --link        Marker-Link, der als "Referenz" gilt und NICHT verwendet wird
                (Default: 'Board')

Output (stdout):
  JSON mit axisDir, axisPoint_mm, residuals, tiltXY_deg, tiltYZ_deg, …
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
from typing import List, Optional, Tuple

import numpy as np


# ── Laden ─────────────────────────────────────────────────────────────────────

def load_poses(path: str) -> dict:
    """Lädt aruco_marker_poses.json – akzeptiert Verzeichnis oder direkten Pfad."""
    if os.path.isdir(path):
        path = os.path.join(path, 'aruco_marker_poses.json')
    if not os.path.exists(path):
        raise FileNotFoundError(f'Nicht gefunden: {path}')
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def fremd_markers(data: dict, ref_link: str = 'Board') -> dict[int, dict]:
    """Gibt fremd-Marker als {marker_id: marker_dict} zurück."""
    return {
        m['marker_id']: m
        for m in data.get('markers', [])
        if m.get('link') != ref_link
    }


def get_points_mm(marker: dict) -> List[np.ndarray]:
    """
    Gibt alle 3D-Punkte eines Markers in mm zurück.
    Bevorzugt corners_m (4 Ecken); Fallback: Zentrum position_mm.
    """
    corners = marker.get('corners_m')
    if corners:
        return [np.array(c, dtype=float) * 1000.0 for c in corners]
    pos = marker.get('position_mm')
    if pos:
        return [np.array(pos, dtype=float)]
    pos_m = marker.get('position_m')
    if pos_m:
        return [np.array(pos_m, dtype=float) * 1000.0]
    return []


# ── Geometrie ─────────────────────────────────────────────────────────────────

def circumcenter_and_normal(
    P1: np.ndarray, P2: np.ndarray, P3: np.ndarray
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], float]:
    """
    Umkreismittelpunkt + Normalenvektor des Dreiecks P1-P2-P3.

    Rückgabe: (circumcenter_mm, normal, cross_length)
    Bei Degeneration (kollinear / identisch): (None, None, cross_length)
    """
    v1 = P2 - P1
    v2 = P3 - P1
    cross = np.cross(v1, v2)
    cross_len = float(np.linalg.norm(cross))

    if cross_len < 1e-6:          # Punkte (fast) kollinear oder identisch
        return None, None, cross_len

    n = cross / cross_len

    # Baryzentrische Gewichte → Umkreismittelpunkt
    a2 = float(np.dot(P2 - P3, P2 - P3))
    b2 = float(np.dot(P1 - P3, P1 - P3))
    c2 = float(np.dot(P1 - P2, P1 - P2))

    w1 = a2 * (b2 + c2 - a2)
    w2 = b2 * (a2 + c2 - b2)
    w3 = c2 * (a2 + b2 - c2)
    w_sum = w1 + w2 + w3

    if abs(w_sum) < 1e-12:        # gleichseitiges / entartetes Dreieck
        return None, None, cross_len

    C = (w1 * P1 + w2 * P2 + w3 * P3) / w_sum
    return C, n, cross_len


def point_line_distance(point: np.ndarray, line_pt: np.ndarray, line_dir: np.ndarray) -> float:
    """Abstand Punkt → Gerade (line_pt + t·line_dir, line_dir normiert)."""
    diff = point - line_pt
    return float(np.linalg.norm(diff - np.dot(diff, line_dir) * line_dir))


# ── Kernberechnung ────────────────────────────────────────────────────────────

def compute_rotation_axis(
    posA: dict,
    posB: dict,
    posC: dict,
    ref_link: str = 'Board',
    min_radius_mm: float = 0.5,
    min_movement_mm: float = 10.0,
) -> dict:
    """
    Berechnet die Rotationsachse aus drei Messungen.

    Parameter:
      min_radius_mm   – Kreisradius unter dem ein Ecken-Triplet als degenerat gilt
      min_movement_mm – Minimale Zentren-Bewegung (max Paarweisabstand A/B/C)
                        Marker die sich weniger bewegen werden ignoriert.
                        Default 10 mm → filtert Board-nahe / fest stehende Marker.
    """
    mA = fremd_markers(posA, ref_link)
    mB = fremd_markers(posB, ref_link)
    mC = fremd_markers(posC, ref_link)

    common_ids = sorted(set(mA) & set(mB) & set(mC))
    if not common_ids:
        return {'ok': False, 'error': 'Keine gemeinsamen fremd-Marker in A+B+C'}

    circumcenters: List[np.ndarray] = []
    normals:       List[np.ndarray] = []
    skipped:       List[dict]       = []
    marker_results: List[dict]      = []

    for mid in common_ids:
        # ── Mindest-Bewegungs-Filter ───────────────────────────────────────────
        # Marker die sich kaum bewegen liefern degenerate Umkreismittelpunkte.
        # Wir vergleichen die Zentren der drei Messungen.
        # Fehlt position_mm in einer Messung (z.B. Einzelkamera-Marker) → überspringen.
        cA_raw = mA[mid].get('position_mm')
        cB_raw = mB[mid].get('position_mm')
        cC_raw = mC[mid].get('position_mm')
        if cA_raw is None or cB_raw is None or cC_raw is None:
            skipped.append({
                'marker_id': mid,
                'reason': 'fehlende position_mm in mindestens einer Messung (z.B. Einzelkamera-Marker)',
            })
            continue
        cA = np.array(cA_raw, dtype=float)
        cB = np.array(cB_raw, dtype=float)
        cC = np.array(cC_raw, dtype=float)
        max_movement = max(
            np.linalg.norm(cB - cA),
            np.linalg.norm(cC - cB),
            np.linalg.norm(cC - cA),
        )
        if max_movement < min_movement_mm:
            skipped.append({
                'marker_id': mid,
                'reason': 'Bewegung zu gering (kein rotierender Marker)',
                'max_movement_mm': round(float(max_movement), 2),
                'threshold_mm': min_movement_mm,
            })
            continue

        pts_a = get_points_mm(mA[mid])
        pts_b = get_points_mm(mB[mid])
        pts_c = get_points_mm(mC[mid])

        n_pts = min(len(pts_a), len(pts_b), len(pts_c))
        ok_pts: List[np.ndarray] = []
        ok_ns:  List[np.ndarray] = []

        for i in range(n_pts):
            P1, P2, P3 = pts_a[i], pts_b[i], pts_c[i]
            C, n, cross_len = circumcenter_and_normal(P1, P2, P3)

            if C is None:
                skipped.append({
                    'marker_id': mid, 'point_idx': i,
                    'reason': 'degenerat (kollinear)',
                    'cross_len': round(cross_len, 6),
                })
                continue

            radius = float(np.linalg.norm(C - P1))
            if radius < min_radius_mm:
                skipped.append({
                    'marker_id': mid, 'point_idx': i,
                    'reason': 'radius zu klein (Bewegung zu gering)',
                    'radius_mm': round(radius, 4),
                })
                continue

            ok_pts.append(C)
            ok_ns.append(n)
            circumcenters.append(C)
            normals.append(n)

        cc_arr = np.array(ok_pts) if ok_pts else np.empty((0, 3))
        marker_results.append({
            'marker_id':       mid,
            'n_points_used':   len(ok_pts),
            'n_points_total':  n_pts,
            'circumcenter_mean_mm': cc_arr.mean(axis=0).round(3).tolist()
            if len(ok_pts) else None,
        })

    if not normals:
        return {
            'ok': False,
            'error': 'Alle Triplets degenerat – Bewegung zu gering oder Marker kollinear.',
            'skipped': skipped,
        }

    normals_arr      = np.array(normals)          # (N, 3)
    circumcenters_arr = np.array(circumcenters)   # (N, 3)

    # ── Achsenrichtung ─────────────────────────────────────────────────────────
    # Alle Normalen auf dieselbe Halbkugel (Vorzeichenanpassung)
    ref = normals_arr[0]
    signs = np.where(normals_arr @ ref >= 0, 1.0, -1.0)
    aligned = normals_arr * signs[:, np.newaxis]

    mean_n   = aligned.mean(axis=0)
    axis_dir = mean_n / np.linalg.norm(mean_n)

    # Streuung der Normalen (Winkelresidum)
    cos_angles = np.clip(aligned @ axis_dir, -1, 1)
    angle_residuals_deg = np.degrees(np.arccos(cos_angles))

    # ── Referenzpunkt ──────────────────────────────────────────────────────────
    axis_point = circumcenters_arr.mean(axis=0)

    # ── Abstandsresiduen ───────────────────────────────────────────────────────
    dist_residuals = np.array([
        point_line_distance(C, axis_point, axis_dir)
        for C in circumcenters_arr
    ])

    # ── Abweichung von Y-Achse [0, 1, 0] in Roboter-Koordinaten ──────────────
    ax, ay, az = axis_dir
    tilt_xy_deg = math.degrees(math.atan2(ax, ay))   # Kippung in XY-Ebene
    tilt_yz_deg = math.degrees(math.atan2(az, ay))   # Kippung in YZ-Ebene

    used_ids = [r['marker_id'] for r in marker_results if r['n_points_used'] > 0]

    return {
        'ok':                   True,
        'axisDir':              axis_dir.round(6).tolist(),
        'axisPoint_mm':         axis_point.round(3).tolist(),
        'tiltXY_deg':           round(tilt_xy_deg, 4),
        'tiltYZ_deg':           round(tilt_yz_deg, 4),
        'numMarkersUsed':       len(used_ids),
        'numMarkersCommon':     len(common_ids),
        'usedMarkerIds':        used_ids,
        'commonMarkerIds':      common_ids,
        'numPoints':            len(normals),
        'residual_dist_mean_mm':  round(float(dist_residuals.mean()), 3),
        'residual_dist_max_mm':   round(float(dist_residuals.max()),  3),
        'residual_dist_mm':       dist_residuals.round(3).tolist(),
        'residual_angle_mean_deg': round(float(angle_residuals_deg.mean()), 4),
        'residual_angle_max_deg':  round(float(angle_residuals_deg.max()),  4),
        'markerResults':        marker_results,
        'skipped':              skipped,
    }


# ── CLI ───────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description='Rotationsachse aus drei Board-Timestamps berechnen'
    )
    parser.add_argument('posA', help='Timestamp A – Verzeichnis oder aruco_marker_poses.json')
    parser.add_argument('posB', help='Timestamp B – Verzeichnis oder aruco_marker_poses.json')
    parser.add_argument('posC', help='Timestamp C – Verzeichnis oder aruco_marker_poses.json')
    parser.add_argument('--output', '-o', metavar='FILE',
                        help='Ergebnis zusätzlich in JSON-Datei speichern')
    parser.add_argument('--pretty', action='store_true',
                        help='Eingerücktes JSON auf stdout')
    parser.add_argument('--link', default='Board',
                        help='Referenz-Link der Board-Marker (Default: Board)')
    parser.add_argument('--min-radius', type=float, default=0.5,
                        help='Min. Kreisradius mm (Default 0.5)')
    parser.add_argument('--min-movement', type=float, default=10.0,
                        help='Min. Marker-Zentrumsbewegung mm (Default 10.0)')
    args = parser.parse_args()

    try:
        data_a = load_poses(args.posA)
        data_b = load_poses(args.posB)
        data_c = load_poses(args.posC)
    except FileNotFoundError as e:
        print(f'Fehler: {e}', file=sys.stderr)
        sys.exit(1)

    result = compute_rotation_axis(data_a, data_b, data_c,
                                   ref_link=args.link,
                                   min_radius_mm=args.min_radius,
                                   min_movement_mm=args.min_movement)

    indent = 2 if args.pretty else None
    out_str = json.dumps(result, indent=indent, ensure_ascii=False)
    print(out_str)

    if args.output:
        with open(args.output, 'w', encoding='utf-8') as f:
            f.write(out_str)
        print(f'Ergebnis gespeichert: {args.output}', file=sys.stderr)

    if result['ok']:
        r = result
        fmt = lambda v: f'{v:+.3f}°'
        print('', file=sys.stderr)
        print('─── Rotationsachse ───────────────────────────────────────', file=sys.stderr)
        print(f"  Richtung:      [{', '.join(f'{v:.4f}' for v in r['axisDir'])}]", file=sys.stderr)
        print(f"  Referenzpunkt: [{', '.join(f'{v:.1f}' for v in r['axisPoint_mm'])}] mm", file=sys.stderr)
        print(f"  Abw. von Y:    XY={fmt(r['tiltXY_deg'])}  YZ={fmt(r['tiltYZ_deg'])}", file=sys.stderr)
        print(f"  Marker: {r['numMarkersUsed']} genutzt / {r['numMarkersCommon']} gemeinsam  "
              f"Punkte: {r['numPoints']}", file=sys.stderr)
        print(f"  Residuen Abstand: ⌀{r['residual_dist_mean_mm']:.2f} mm  "
              f"max={r['residual_dist_max_mm']:.2f} mm", file=sys.stderr)
        print(f"  Residuen Winkel:  ⌀{r['residual_angle_mean_deg']:.3f}°  "
              f"max={r['residual_angle_max_deg']:.3f}°", file=sys.stderr)
        print('──────────────────────────────────────────────────────────', file=sys.stderr)
    else:
        print(f"Fehler: {result['error']}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()