#!/usr/bin/env python3
"""Planar point-to-point ICP: three small, synthetic teaching examples.

Python 3.10+; standard library only. Coordinates are meters; angles are radians.
The solver receives point arrays and an initial pose, never a reference answer.
"""

import argparse
from dataclasses import dataclass
import json
import math
from pathlib import Path
import sys


Point = tuple[float, float]
MAX_ITERATIONS = 30
STEP_TOLERANCE = 1e-10
CASE_NAMES = ("nearby", "bad-start", "symmetric")


@dataclass(frozen=True)
class Pose:
    """Source-to-target transform: R(theta) p + (x, y)."""

    x: float = 0.0
    y: float = 0.0
    theta: float = 0.0


@dataclass(frozen=True)
class Snapshot:
    iteration: int
    transform: Pose
    points: list[Point]
    matches: list[int]
    rmse: float
    status: str
    translation_step: float | None = None
    rotation_step: float | None = None


def check_points(points: list[Point]) -> None:
    if len(points) < 2 or any(
        len(point) != 2 or any(not math.isfinite(value) for value in point)
        for point in points
    ):
        raise ValueError("Use at least two finite planar points.")


def wrap_angle(angle: float) -> float:
    return math.atan2(math.sin(angle), math.cos(angle))


def apply_transform(point: Point, transform: Pose) -> Point:
    c, s = math.cos(transform.theta), math.sin(transform.theta)
    x, y = point
    return c * x - s * y + transform.x, s * x + c * y + transform.y


def compose_transforms(outer: Pose, inner: Pose) -> Pose:
    """Return outer(inner(point)); rotate the old translation before adding."""
    x, y = apply_transform((inner.x, inner.y), outer)
    return Pose(x, y, wrap_angle(outer.theta + inner.theta))


def fit_rigid_2d(source: list[Point], target: list[Point]) -> Pose | None:
    """Fit fixed pairs; return None when rotation is numerically undetermined."""
    check_points(source)
    check_points(target)
    if len(source) != len(target):
        raise ValueError("A fixed-pair fit needs equally sized point arrays.")
    n = len(source)
    p = tuple(sum(point[j] / n for point in source) for j in (0, 1))
    q = tuple(sum(point[j] / n for point in target) for j in (0, 1))
    a = b = source_spread = target_spread = 0.0
    for point, matched in zip(source, target):
        px, py = point[0] - p[0], point[1] - p[1]
        qx, qy = matched[0] - q[0], matched[1] - q[1]
        a += px * qx + py * qy
        b += px * qy - py * qx
        source_spread += px * px + py * py
        target_spread += qx * qx + qy * qy
    scale = math.sqrt(source_spread) * math.sqrt(target_spread)
    if not all(math.isfinite(value) for value in (a, b, scale)):
        raise ValueError("Point spread exceeds the numeric range.")
    if scale == 0 or math.hypot(a, b) <= 64 * sys.float_info.epsilon * scale:
        return None
    theta = math.atan2(b, a)
    rotated = apply_transform(p, Pose(theta=theta))
    return Pose(q[0] - rotated[0], q[1] - rotated[1], theta)


def nearest(points: list[Point], target: list[Point]) -> tuple[list[int], float]:
    """Directed matches: each source chooses one target; first-index ties win."""
    matches, squared_sum = [], 0.0
    for x, y in points:
        best, best_squared = 0, math.inf
        for index, (qx, qy) in enumerate(target):
            dx, dy = x - qx, y - qy
            squared = dx * dx + dy * dy
            if squared < best_squared:
                best, best_squared = index, squared
        matches.append(best)
        squared_sum += best_squared
    rmse = math.sqrt(squared_sum / len(points))
    if not math.isfinite(rmse):
        raise ValueError("Point distances exceed the numeric range.")
    return matches, rmse


def run_icp(
    source: list[Point], target: list[Point], initial: Pose = Pose(),
    max_iterations: int = MAX_ITERATIONS,
) -> list[Snapshot]:
    """Match, fit current points, left-compose, and rematch; no truth input."""
    check_points(source)
    check_points(target)
    if not all(math.isfinite(value) for value in (initial.x, initial.y, initial.theta)):
        raise ValueError("Use a finite rigid transform.")
    if type(max_iterations) is not int or not 1 <= max_iterations <= 100:
        raise ValueError("Use an iteration budget from 1 to 100.")
    points = [apply_transform(point, initial) for point in source]
    matches, rmse = nearest(points, target)
    history = [Snapshot(0, initial, points, matches, rmse, "Ready")]
    for iteration in range(1, max_iterations + 1):
        previous = history[-1]
        delta = fit_rigid_2d(previous.points, [target[i] for i in previous.matches])
        if delta is None:
            history.append(Snapshot(iteration, previous.transform, previous.points,
                                    previous.matches, previous.rmse, "Rotation undetermined"))
            break
        transform = compose_transforms(delta, previous.transform)
        points = [apply_transform(point, transform) for point in source]
        matches, rmse = nearest(points, target)
        translation_step = math.hypot(delta.x, delta.y)
        rotation_step = abs(wrap_angle(delta.theta))
        if translation_step < STEP_TOLERANCE and rotation_step < STEP_TOLERANCE:
            status = "Step tolerance reached"
        elif iteration == max_iterations:
            status = "Iteration limit"
        else:
            status = "Iterating"
        history.append(Snapshot(iteration, transform, points, matches, rmse, status,
                                translation_step, rotation_step))
        if status != "Iterating":
            break
    return history


def load_examples() -> dict:
    """Read stored source/target coordinates; never generate a target here."""
    with Path(__file__).with_name("points.json").open(encoding="utf-8") as stream:
        data = json.load(stream)
    for name in ("asymmetric", "symmetric"):
        for role in ("source", "target"):
            points = data["datasets"][name][role]
            check_points(points)
    for name in CASE_NAMES:
        case = data["cases"][name]
        if case["dataset"] not in ("asymmetric", "symmetric"):
            raise ValueError(f"Unknown dataset for {name}.")
        pose_from_dict(case["initial"])
    pose_from_dict(data["evaluation"]["reference_transform"])
    return data


def pose_from_dict(value: dict) -> Pose:
    pose = Pose(value["x"], value["y"], value["theta"])
    if not all(math.isfinite(item) for item in (pose.x, pose.y, pose.theta)):
        raise ValueError("Use a finite rigid transform.")
    return pose


def case_history(data: dict, name: str) -> list[Snapshot]:
    case = data["cases"][name]
    dataset = data["datasets"][case["dataset"]]
    return run_icp(dataset["source"], dataset["target"], pose_from_dict(case["initial"]))


def pose_errors(estimate: Pose, reference: Pose) -> tuple[float, float]:
    """Evaluation only; the solver never calls this function."""
    translation = math.hypot(estimate.x - reference.x, estimate.y - reference.y)
    rotation = abs(math.degrees(wrap_angle(estimate.theta - reference.theta)))
    return translation, rotation


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(f"Check failed: {message}")


def check_examples(data: dict, histories: dict[str, list[Snapshot]]) -> None:
    """Check independent hand calculations and fixed reference benchmarks."""
    triangle = [(0, 0), (1, 0), (0, 1)]
    paired = [(1, 2), (1, 3), (0, 2)]
    fit = fit_rigid_2d(triangle, paired)
    require(fit is not None, "fixed-pair triangle must have a determined rotation")
    require(math.hypot(fit.x - 1, fit.y - 2) < 1e-12
            and abs(fit.theta - math.pi / 2) < 1e-12,
            "fixed-pair triangle must give 90 degrees and (1, 2) m")
    combined = compose_transforms(Pose(2, -1, math.pi / 2), Pose(1, 0, math.pi / 2))
    require(math.hypot(combined.x - 2, combined.y) < 1e-12
            and abs(abs(combined.theta) - math.pi) < 1e-12,
            "composition must rotate the old translation")
    require(nearest([(0, 0)], [(-1, 0), (1, 0)])[0] == [0],
            "an exact nearest-neighbor tie must keep target index zero")
    require(fit_rigid_2d([(0, 0), (0, 0)], [(1, 1), (1, 1)]) is None,
            "collapsed pairs cannot determine rotation")
    reference = pose_from_dict(data["evaluation"]["reference_transform"])
    require(pose_errors(reference, Pose(.6, -.2, math.radians(20))) == (0.0, 0.0),
            "evaluation reference must stay at 20 degrees and (0.6, -0.2) m")
    # Values are fixed benchmarks, not calculated from the current solver run.
    benchmarks = {
        "nearby": (0.562084, 0.432199, 0.000000, 0.000000, 0.000000, 4),
        "bad-start": (0.695486, 0.466319, 0.366997, 0.405761, 128.866330, 5),
        "symmetric": (0.000000, None, 0.000000, 0.000000, 180.000000, 1),
    }
    for name, history in histories.items():
        initial_rmse, first_rmse, final_rmse, translation_error, rotation_error, steps = benchmarks[name]
        final = history[-1]
        measured_translation, measured_rotation = pose_errors(final.transform, reference)
        comparisons = [(history[0].rmse, initial_rmse), (final.rmse, final_rmse),
                       (measured_translation, translation_error), (measured_rotation, rotation_error)]
        if first_rmse is not None:
            comparisons.append((history[1].rmse, first_rmse))
        require(all(abs(actual - expected) < 1e-6 for actual, expected in comparisons),
                f"{name} must reproduce its stored numerical benchmarks")
        require(final.iteration == steps and final.status == "Step tolerance reached",
                f"{name} must stop after the expected tiny update")
        require(all(b.rmse <= a.rmse + 1e-12 for a, b in zip(history, history[1:])),
                f"{name} nearest-point RMSE must not increase beyond rounding")
    good, ambiguous = histories["nearby"][-1], histories["symmetric"][-1]
    require(good.rmse < 1e-12 and max(pose_errors(good.transform, reference)) < 1e-9,
            "nearby must recover the reference pose")
    require(ambiguous.rmse < 1e-12 and pose_errors(ambiguous.transform, reference)[1] > 179.999,
            "symmetric must retain the zero-residual, wrong-rotation example")
    bad = histories["bad-start"][-1].transform
    require(abs(bad.x - .539248) < 1e-6 and abs(bad.y + .601187) < 1e-6
            and abs(math.degrees(bad.theta) - 148.866330) < 1e-6,
            "bad-start must retain its wrong final pose")


def print_case(name: str, history: list[Snapshot], reference: Pose, trace: bool) -> None:
    final = history[-1]
    translation_error, rotation_error = pose_errors(final.transform, reference)
    print(f"[{name}]")
    print(f"Initial nearest-point RMSE: {history[0].rmse:.6f} m")
    print(f"After one update: {history[1].rmse:.6f} m")
    print(f"Final pose: {math.degrees(final.transform.theta):.6f} deg, "
          f"({final.transform.x:.6f}, {final.transform.y:.6f}) m")
    print(f"Final nearest-point RMSE: {final.rmse:.6f} m")
    print(f"Reference pose errors: {translation_error:.6f} m, {rotation_error:.6f} deg")
    print(f"Stop: {final.status} (iteration {final.iteration})")
    if trace:
        print("iteration,theta_deg,x_m,y_m,rmse_m,distinct_targets,status")
        for item in history:
            pose = item.transform
            print(f"{item.iteration},{math.degrees(pose.theta):.9f},{pose.x:.9f},"
                  f"{pose.y:.9f},{item.rmse:.9f},{len(set(item.matches))},{item.status}")
    print()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--case", choices=CASE_NAMES, help="print one case (default: all three)")
    parser.add_argument("--trace", action="store_true", help="print the iteration history")
    parser.add_argument("--check", action="store_true", help="verify all three cases and hand-worked checks")
    args = parser.parse_args()
    try:
        data = load_examples()
        names = CASE_NAMES if args.check or args.case is None else (args.case,)
        histories = {name: case_history(data, name) for name in names}
        reference = pose_from_dict(data["evaluation"]["reference_transform"])
        for name in (CASE_NAMES if args.case is None else (args.case,)):
            print_case(name, histories[name], reference, args.trace)
        if args.check:
            check_examples(data, histories)
            print("Checks passed: fixed-pair fit, composition, ties, degeneracy, and all three cases.")
        return 0
    except (OSError, ValueError, TypeError, KeyError, OverflowError) as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())