#!/usr/bin/env python3
"""Evaluate timestamped planar poses using only the Python standard library."""
import argparse
from bisect import bisect_left, bisect_right
import csv
import hashlib
import json
from math import atan2, cos, degrees, fsum, hypot, isfinite, sin, sqrt
from pathlib import Path

ROOT = Path(__file__).resolve().parent
TIME_EPSILON = 1e-12
CASES = (
    ("interior", "interior-error.csv", "none", 0.0),
    ("heading", "heading-drift.csv", "none", 0.0),
    ("frame-raw", "frame-offset.csv", "none", 0.0),
    ("frame-rigid", "frame-offset.csv", "rigid", 0.0),
    ("scale-rigid", "scale-error.csv", "rigid", 0.0),
    ("scale-similarity", "scale-error.csv", "similarity", 0.0),
    ("clock-raw", "clock-offset.csv", "none", 0.0),
    ("clock-rigid", "clock-offset.csv", "rigid", 0.0),
    ("clock-corrected", "clock-offset.csv", "none", -0.4),
)


def validate_poses(poses):
    if len(poses) < 2:
        raise ValueError("Supply at least two poses in each trajectory.")
    previous = None
    for index, pose in enumerate(poses):
        if set(pose) != {"t", "x", "y", "yaw"}:
            raise ValueError(f"Pose {index}: expected t,x,y,yaw fields.")
        if not all(isinstance(v, (int, float)) and isfinite(v) for v in pose.values()):
            raise ValueError(f"Pose {index}: all values must be finite numbers.")
        if previous is not None and pose["t"] <= previous:
            raise ValueError("Timestamps must be strictly increasing; duplicates are invalid.")
        previous = pose["t"]


def read_poses(path):
    with Path(path).open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        if reader.fieldnames != ["t", "x", "y", "yaw"]:
            raise ValueError(f"{path}: CSV header must be exactly t,x,y,yaw.")
        poses = []
        for row in reader:
            if None in row or any(value is None for value in row.values()):
                raise ValueError(f"{path}: row {reader.line_num} has the wrong column count.")
            try:
                poses.append({key: float(value) for key, value in row.items()})
            except ValueError as error:
                raise ValueError(f"{path}: row {reader.line_num} must contain numbers.") from error
    validate_poses(poses)
    return poses


def associate(reference, estimate, offset=0.0, max_difference=0.15):
    """Global greedy, one-to-one matching; offset is ADDED to estimate time."""
    validate_poses(reference)
    validate_poses(estimate)
    if not isfinite(offset) or not isfinite(max_difference) or max_difference < 0:
        raise ValueError("Offset must be finite and max difference finite and nonnegative.")
    times = [p["t"] + offset for p in estimate]
    if not all(isfinite(t) for t in times):
        raise ValueError("Corrected estimate timestamps must remain finite.")
    candidates = []
    for i, pose in enumerate(reference):
        lower = bisect_left(times, pose["t"] - max_difference - TIME_EPSILON)
        upper = bisect_right(times, pose["t"] + max_difference + TIME_EPSILON)
        for j in range(lower, upper):
            difference = abs(times[j] - pose["t"])
            if difference <= max_difference + TIME_EPSILON:
                candidates.append((difference, i, j))
    used_reference, used_estimate, matches = set(), set(), []
    for _, i, j in sorted(candidates):
        if i not in used_reference and j not in used_estimate:
            matches.append((i, j))
            used_reference.add(i)
            used_estimate.add(j)
    return sorted(matches)


def wrap(angle):
    return atan2(sin(angle), cos(angle))


def fit_alignment(reference, estimate, mode):
    """Equal-weight position fit from matched estimate into reference world."""
    if mode not in ("none", "rigid", "similarity"):
        raise ValueError("Alignment must be none, rigid, or similarity.")
    if len(reference) != len(estimate) or len(reference) < 2:
        raise ValueError("Alignment needs at least two corresponding poses.")
    if mode == "none":
        return {"mode": mode, "scale": 1.0, "yaw": 0.0, "tx": 0.0, "ty": 0.0}
    means = [[fsum(p[k] for p in poses) / len(poses) for k in ("x", "y")]
             for poses in (estimate, reference)]
    (px, py), (qx, qy) = means
    centered = [(p["x"] - px, p["y"] - py, q["x"] - qx, q["y"] - qy)
                for p, q in zip(estimate, reference)]
    dot = fsum(ux * vx + uy * vy for ux, uy, vx, vy in centered)
    cross = fsum(ux * vy - uy * vx for ux, uy, vx, vy in centered)
    spread = fsum(ux * ux + uy * uy for ux, uy, _, _ in centered)
    ref_spread = fsum(vx * vx + vy * vy for _, _, vx, vy in centered)
    correlation = hypot(dot, cross)
    if min(spread, ref_spread) <= 1e-24 or correlation <= 1e-12 * sqrt(spread * ref_spread):
        raise ValueError("Matched positions do not determine a stable alignment rotation.")
    yaw = atan2(cross, dot)
    scale = correlation / spread if mode == "similarity" else 1.0
    return {"mode": mode, "scale": scale, "yaw": yaw,
            "tx": qx - scale * (cos(yaw) * px - sin(yaw) * py),
            "ty": qy - scale * (sin(yaw) * px + cos(yaw) * py)}


def apply_alignment(pose, alignment):
    c, s, scale = cos(alignment["yaw"]), sin(alignment["yaw"]), alignment["scale"]
    return {"t": pose["t"], "x": scale * (c * pose["x"] - s * pose["y"]) + alignment["tx"],
            "y": scale * (s * pose["x"] + c * pose["y"]) + alignment["ty"],
            "yaw": wrap(pose["yaw"] + alignment["yaw"])}


def relative_pose(first, second):
    dx, dy = second["x"] - first["x"], second["y"] - first["y"]
    c, s = cos(first["yaw"]), sin(first["yaw"])
    return {"x": c * dx + s * dy, "y": -s * dx + c * dy,
            "yaw": wrap(second["yaw"] - first["yaw"])}


def rms(values):
    return sqrt(fsum(value * value for value in values) / len(values))


def evaluate(reference, estimate, *, offset=0.0, max_difference=0.15,
             alignment="none", lag=1.0, lag_tolerance=0.01):
    if not isfinite(lag) or lag <= 0 or not isfinite(lag_tolerance) or lag_tolerance < 0:
        raise ValueError("Lag must be finite and positive; lag tolerance finite and nonnegative.")
    indices = associate(reference, estimate, offset, max_difference)
    if len(indices) < 2:
        raise ValueError("Fewer than two timestamp matches; check offset, tolerance, and coverage.")
    if any(right[1] <= left[1] for left, right in zip(indices, indices[1:])):
        raise ValueError("Matches reverse estimate time; reduce the matching tolerance or check clocks.")
    refs = [reference[i] for i, _ in indices]
    estimates = [estimate[j] for _, j in indices]
    transform = fit_alignment(refs, estimates, alignment)
    aligned = [apply_alignment(pose, transform) for pose in estimates]
    matches = []
    for match_index, ((i, j), q, p) in enumerate(zip(indices, refs, aligned)):
        matches.append({"match_index": match_index, "reference_index": i, "estimate_index": j,
                        "reference_t": q["t"], "estimate_t": estimate[j]["t"],
                        "corrected_estimate_t": estimate[j]["t"] + offset,
                        "time_residual_s": estimate[j]["t"] + offset - q["t"],
                        "reference_x": q["x"], "reference_y": q["y"], "reference_yaw": q["yaw"],
                        "aligned_x": p["x"], "aligned_y": p["y"], "aligned_yaw": p["yaw"],
                        "position_error_m": hypot(p["x"] - q["x"], p["y"] - q["y"])})
    pairs = []
    times = [pose["t"] for pose in refs]
    for i, start in enumerate(times[:-1]):
        target = start + lag
        insertion = bisect_left(times, target, lo=i + 1)
        candidates = [j for j in (insertion - 1, insertion) if i < j < len(times)]
        j = min(candidates, key=lambda index: (abs(times[index] - target), index))
        interval = times[j] - start
        if abs(interval - lag) > lag_tolerance + TIME_EPSILON:
            continue
        reference_motion = relative_pose(refs[i], refs[j])
        estimate_motion = relative_pose(aligned[i], aligned[j])
        error = relative_pose(reference_motion, estimate_motion)
        pairs.append({"start_match": i, "end_match": j,
                      "reference_start_index": indices[i][0], "reference_end_index": indices[j][0],
                      "estimate_start_index": indices[i][1], "estimate_end_index": indices[j][1],
                      "reference_interval_s": interval,
                      "estimate_interval_s": estimates[j]["t"] - estimates[i]["t"],
                      "translation_error_m": hypot(error["x"], error["y"]),
                      "rotation_error_rad": abs(error["yaw"])})
    if not pairs:
        raise ValueError("No eligible RPE pairs; check the lag, lag tolerance, and matched time span.")
    used_ref, used_est = {i for i, _ in indices}, {j for _, j in indices}
    residuals = [abs(match["time_residual_s"]) for match in matches]
    return {
        "protocol": {"version": "planar-trajectory-v1", "position_units": "m", "time_units": "s",
                     "angle_units": "rad", "pose_convention": "body-to-world",
                     "estimate_time_offset_s": offset, "max_difference_s": max_difference,
                     "boundary_epsilon_s": TIME_EPSILON,
                     "association": "global greedy by absolute residual, then reference index, then estimate index",
                     "interpolation": False, "requested_lag_s": lag, "lag_tolerance_s": lag_tolerance,
                     "rpe_pair_selection": "nearest later matched reference time to start + lag; ties choose earlier",
                     "alignment_fit": "not fitted" if alignment == "none" else "all matched positions; equal weight",
                     "rpe_uses_aligned_poses": True},
        "association": {"reference_count": len(reference), "estimate_count": len(estimate),
                        "matched_count": len(matches), "reference_coverage": len(matches) / len(reference),
                        "estimate_coverage": len(matches) / len(estimate),
                        "unmatched_reference_indices": [i for i in range(len(reference)) if i not in used_ref],
                        "unmatched_estimate_indices": [i for i in range(len(estimate)) if i not in used_est],
                        "mean_abs_time_residual_s": fsum(residuals) / len(residuals),
                        "max_abs_time_residual_s": max(residuals)},
        "alignment": transform,
        "metrics": {"position_rmse_m": rms([m["position_error_m"] for m in matches]),
                    "last_matched_position_error_m": matches[-1]["position_error_m"],
                    "rpe_pair_count": len(pairs),
                    "rpe_translation_rmse_m": rms([p["translation_error_m"] for p in pairs]),
                    "rpe_rotation_rmse_deg": degrees(rms([p["rotation_error_rad"] for p in pairs])),
                    "reference_interval_min_s": min(p["reference_interval_s"] for p in pairs),
                    "reference_interval_max_s": max(p["reference_interval_s"] for p in pairs),
                    "estimate_interval_min_s": min(p["estimate_interval_s"] for p in pairs),
                    "estimate_interval_max_s": max(p["estimate_interval_s"] for p in pairs)},
        "matches": matches, "rpe_pairs": pairs,
    }


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def evaluate_files(reference, estimate, **options):
    result = evaluate(read_poses(reference), read_poses(estimate), **options)
    result["inputs"] = {"reference": {"name": Path(reference).name, "sha256": sha256(reference)},
                        "estimate": {"name": Path(estimate).name, "sha256": sha256(estimate)}}
    return result


def write_json(path, value):
    path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8")


def write_report(output, result):
    if output.exists() and (not output.is_dir() or any(output.iterdir())):
        raise ValueError(f"Output must be a new or empty directory: {output}")
    output.mkdir(parents=True, exist_ok=True)
    write_json(output / "report.json", result)
    for name, rows in (("matches.csv", result["matches"]), ("rpe.csv", result["rpe_pairs"])):
        with (output / name).open("w", newline="", encoding="utf-8") as handle:
            writer = csv.DictWriter(handle, fieldnames=list(rows[0]), lineterminator="\n")
            writer.writeheader()
            writer.writerows(rows)


def summary_line(name, report):
    a, m = report["association"], report["metrics"]
    return (f"{name}: matches={a['matched_count']}/{a['reference_count']} "
            f"ATE={m['position_rmse_m']:.6f}m scale={report['alignment']['scale']:.6f} "
            f"RPE({report['protocol']['requested_lag_s']:g}s)={m['rpe_translation_rmse_m']:.6f}m/"
            f"{m['rpe_rotation_rmse_deg']:.6f}deg pairs={m['rpe_pair_count']}")


def demo(output=None):
    if output is not None and output.exists() and (not output.is_dir() or any(output.iterdir())):
        raise ValueError(f"Output must be a new or empty directory: {output}")
    reports = {}
    for name, filename, mode, offset in CASES:
        result = evaluate_files(ROOT / "data/reference.csv", ROOT / "data" / filename,
                                alignment=mode, offset=offset)
        reports[name] = result
        if output is not None:
            write_report(output / name, result)
    return reports


def compare_values(expected, actual, location="report"):
    if isinstance(expected, dict):
        if not isinstance(actual, dict) or expected.keys() != actual.keys():
            raise ValueError(f"{location}: report keys differ.")
        for key in expected:
            compare_values(expected[key], actual[key], location + "." + key)
    elif isinstance(expected, list):
        if not isinstance(actual, list) or len(expected) != len(actual):
            raise ValueError(f"{location}: report list differs.")
        for index, (left, right) in enumerate(zip(expected, actual)):
            compare_values(left, right, location + f"[{index}]")
    elif isinstance(expected, (float, int)) and not isinstance(expected, bool):
        if not isinstance(actual, (float, int)) or not isfinite(actual) or abs(expected - actual) > 1e-9:
            raise ValueError(f"{location}: expected {expected}, got {actual}.")
    elif expected != actual:
        raise ValueError(f"{location}: expected {expected!r}, got {actual!r}.")


def check():
    expected = json.loads((ROOT / "expected-results.json").read_text(encoding="utf-8"))
    reports = demo()
    # Reference summaries retain input hashes, association details, fit, and metrics.
    actual = {name: {key: value for key, value in report.items() if key not in ("matches", "rpe_pairs")}
              for name, report in reports.items()}
    compare_values(expected, actual)
    output = "\n".join(summary_line(name, report) for name, report in reports.items()) + "\n"
    if output != (ROOT / "expected-output.txt").read_text(encoding="utf-8"):
        raise ValueError("Formatted output differs from expected-output.txt.")
    print(output, end="")
    print("Checked CSV hashes and all nine reference evaluations (numeric tolerance 1e-9).")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("check", help="Verify the bundled CSVs and nine reference evaluations.")
    demo_parser = commands.add_parser("demo", help="Evaluate all nine controlled examples.")
    demo_parser.add_argument("--output", type=Path, default=Path("run"))
    evaluator = commands.add_parser("evaluate", help="Evaluate your own timestamped planar CSVs.")
    evaluator.add_argument("--reference", type=Path, required=True)
    evaluator.add_argument("--estimate", type=Path, required=True)
    evaluator.add_argument("--output", type=Path, required=True)
    evaluator.add_argument("--alignment", choices=("none", "rigid", "similarity"), default="none")
    evaluator.add_argument("--offset", type=float, default=0.0, help="Seconds ADDED to estimate timestamps.")
    evaluator.add_argument("--max-difference", type=float, default=0.15)
    evaluator.add_argument("--lag", type=float, default=1.0)
    evaluator.add_argument("--lag-tolerance", type=float, default=0.01)
    args = parser.parse_args()
    try:
        if args.command == "check":
            check()
        elif args.command == "demo":
            reports = demo(args.output)
            for name, report in reports.items():
                print(summary_line(name, report))
        else:
            result = evaluate_files(args.reference, args.estimate, alignment=args.alignment,
                                    offset=args.offset, max_difference=args.max_difference,
                                    lag=args.lag, lag_tolerance=args.lag_tolerance)
            write_report(args.output, result)
            print(summary_line("evaluation", result))
    except (ValueError, OSError, OverflowError) as error:
        parser.exit(1, f"Error: {error}\n")


if __name__ == "__main__":
    main()